import math
from character import CharacterPlan
from game import Difficulty, Roll, d
from creature import CreaturePlan
from actions import Attack
from items import Weapon

from collections import defaultdict

import matplotlib.pyplot as plt
from matplotlib.ticker import AutoLocator, AutoMinorLocator, MultipleLocator

import pb2e

plt.style.use('dark_background')

def treat_wounds(check_res, heal_bonus=0):
    if check_res <= Roll.Fumble:
        return -d(8)
    if check_res == Roll.Fail:
        return 0
    if check_res == Roll.Success:
        return 2*d(8) + heal_bonus
    if check_res == Roll.Crit:
        return 4*d(8) + heal_bonus

def strike(check_res, dmg, crit_dmg=0):
    if check_res == Roll.Success:
        return dmg
    if check_res == Roll.Crit:
        return 2*dmg + crit_dmg
    return 0


def show_stat(name, dist):
    print(name)
    values = list(dist.probabilities.keys())
    mn = min(values)
    mx = max(values)
    print("min", mn, dist.probabilities[mn])
    print("max", mx, dist.probabilities[mx])
    print("Exp", dist.expected())
    print("Std", dist.standard_deviation())
    print("<0", dist.accumulate_prob(lambda x: x < 0))
    print(">0", dist.accumulate_prob(lambda x: x > 0))
    print("")
    assert abs(dist.accumulate_prob(lambda x: True) - 1) < math.pow(10, -10)



def main():
    treat_options = [
        {"dc": 15, "heal_bonus": 0},
        {"dc": 20, "heal_bonus": 15}
    ]
    bonus = 10
    show_stat("Assurance 15", 2*d(8))
    for opt in treat_options:
        dc = opt["dc"]
        treat = d(20).apply(Roll.check, dc, bonus).apply(treat_wounds, opt["heal_bonus"])
        risky = d(20).apply(Roll.check, dc, bonus+2).apply(lambda x: Roll.Crit if x == Roll.Success else x)
        risky = risky.apply(treat_wounds, opt["heal_bonus"]) + (-d(8))

        show_stat("Treat Wounds " + str(dc), treat)
        show_stat("Risky Surgery " + str(dc), risky)

    treat = d(20).apply(Roll.check, 20, bonus+1).apply(treat_wounds, 15)
    show_stat("Treat Wounds 20 +1", treat)


def test3():
    # char_js = pb2e.download_character(pfb_id=160180)
    char_js = pb2e.open_character("quizrel2")
    char_plan = CharacterPlan(char_js)
    creature_plan = CreaturePlan(default_difficulty=Difficulty.Moderate)
    # creature_plan["reflex"] = Difficulty.High
    # creature_plan["ac"] = Difficulty.High

    def calc_damage(level):
        char = char_plan.build(level)
        creature = creature_plan.build(level)
        creature["ac"] -= 2
        falcata = Weapon("falcata", "martial", "sword", 8, fatal=12)
        sword = Weapon("sword", "martial", "sword", 8)
        rapier = Weapon("rapier", "martial", "sword", 6, finesse=True, deadly=8)
        sword2 = Weapon("sword_d6", "martial", "sword", 6)
        char.weapons = [falcata, sword, rapier, sword2]
        char.update_fundamental_runes()

        # print(char.toJSON())
        # attack = Attack(char, creature, falcata)
        # print(Roll.short_str(attack.roll()))
        # print(f"Attack +{attack.bonus()} vs AC {attack.target.ac()} dmg +{attack.damage_bonus()}")
        # show_stat("roll", attack.roll())
        # show_stat("falcata", attack.strike())

        res = dict()
        for weapon in char.weapons:
            attack = Attack(char, creature, weapon)
            res[weapon.name] = attack.strike().expected()
        # res["hp"] = creature["hp"]
        return res

    # calc_damage(20)

    levelup_view("Weapons Comparison", "Expected Damage", calc_damage)



def levelup_view(title, ylabel, fn):
    # Example data
    levels = list(range(1, 21))
    table = defaultdict(list)
    for level in levels:
        res = fn(level)
        for k,v in res.items():
            table[k].append(v)

    fig, ax = plt.subplots()
    fig.set_size_inches(18.5, 10.5)

    # Plot multiple lines on the same plane
    for label, values in table.items():
        plt.plot(levels, values, label=label)

    # Add labels and legend
    plt.xlabel("Level")
    plt.ylabel(ylabel)
    plt.title(title)
    ax.grid(True, linestyle='--', alpha=0.4)
    ax.xaxis.set_major_locator(MultipleLocator(1))
    ax.yaxis.set_major_locator(AutoLocator())
    ax.yaxis.set_minor_locator(AutoMinorLocator())
    plt.legend()
    plt.show()


if __name__ == '__main__':
    test3()