import math
from distribution import Distribution


def d(edge_count):
    res = Distribution()
    if edge_count == 0:
        res.insert(0, 1)
        return res
    prob = 1/edge_count
    for val in range(1, edge_count+1):
        res.insert(val, prob)
    return res


class Roll:
    Fumble = -1
    Fail = 0
    Success = 1
    Crit = 2

    @staticmethod
    def get_natural_success_mod(value):
        if value <= 1:
            return -1
        if value >= 20:
            return 1
        return 0

    @staticmethod
    def check(value, dc, bonus=0):
        res = 0
        rolled = value + bonus
        if rolled >= dc + 10:
            res = Roll.Crit
        elif rolled >= dc:
            res = Roll.Success
        elif rolled <= dc - 10:
            res = Roll.Fumble
        res += Roll.get_natural_success_mod(value)
        return Roll.clamp(res)

    @staticmethod
    def clamp(value):
        return max(Roll.Fumble, min(Roll.Crit, value))


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)


coin = d(2)


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 main2():
    ac = 18
    bonus = +8
    dmg = d(6)+4
    attack = d(20).apply(Roll.check, ac, bonus)
    show_stat("attack", attack)
    total_dmg = attack.apply(strike, dmg+1)
    show_stat("dogslicer", total_dmg)
    total_dmg = attack.apply(strike, dmg, d(8))
    show_stat("rapier", total_dmg)


if __name__ == '__main__':
    main2()