path: root/main.py

import math

from numpy._core import multiarray
from character import CharacterPlan
from distribution import Distribution
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 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)

class StatsTest:
    def __init__(self, difficulty = Difficulty.Moderate) -> None:
        # self.char_js = pb2e.download_character(pfb_id=160180)
        self.char_js = pb2e.open_character("quizrel2")
        self.char_plan = CharacterPlan(self.char_js)
        self.creature_plan = CreaturePlan(default_difficulty=difficulty)
        # self.creature_plan["reflex"] = Difficulty.High
        # self.creature_plan["ac"] = Difficulty.High
        self.char = self.char_plan.build(1)

    def add_weapons(self):
        wps = [Weapon("falcata", "martial", "sword", 8, fatal=12),
        Weapon("trident", "martial", "spear", 8),
        Weapon("pick", "martial", "pick", 6, fatal=10),
        Weapon("rapier", "martial", "sword", 6, finesse=True, deadly=8),
        Weapon("urumi", "martial", "flail", 6, deadly=10),
        Weapon("tricky pick", "martial", "pick", 6, fatal=10, backstabber=True),
        Weapon("shortsword", "simple", "sword", 6, agile=True),
        Weapon("gauntlet", "simple", "brawling", 4, agile=True),
        Weapon("light pick", "martial", "pick", 6, agile=True, fatal=8)
        ]
        self.char.weapons.extend(wps)

    def get_weapon(self, name):
        for wp in self.char.weapons:
            if wp.name == name:
                return wp
        raise KeyError

    def set_level(self, level):
        self.char = self.char_plan.build(level)
        self.creature = self.creature_plan.build(level)
        self.add_weapons()
        self.char.update_fundamental_runes()

def test4():
    t = StatsTest()
    t.set_level(5)
    urumi = t.get_weapon("urumi")
    attack = Attack(t.char, t.creature, urumi)
    print("Attack: ", Roll.short_str(attack.roll()))
    print(f"Attack +{attack.bonus()} vs AC {attack.target.ac()} dmg +{attack.damage_bonus()}")
    crit_chance = attack.roll().probabilities[Roll.Crit]
    print(t.char.classDC() , t.creature["reflex"])
    prone_save = d(20).apply(Roll.check, t.char.classDC(), t.creature["reflex"])
    print("Prone save: ", Roll.short_str(prone_save))
    prone_chace = prone_save.accumulate_prob(lambda x: x <= Roll.Fail)
    print(f"Crit: {crit_chance}; Prone: {prone_chace}; Total Prone Chance: {crit_chance*prone_chace}")

def make_check(name, bonus, dc):
    roll = d(20).apply2(lambda v, bonus: Roll.check(v, dc, bonus), bonus)
    if isinstance(bonus, Distribution):
        bonus = bonus.expected()
    print(f"{name} {bonus:+} vs DC {dc}:", 
          Roll.short_str(roll), f"Total: {10+bonus-dc:+}")
    return roll

def demoralize_test():
    def bon_mot_pen(check):
        if check == Roll.Success:
            return -2
        if check == Roll.Crit:
            return -3
        return 0
    t = StatsTest()
    t.creature_plan["will"] = Difficulty.Moderate
    t.set_level(5)
    t.char.features["diplomacy"] = 2
    t.char.features["intimidation"] = 4
    intimid_bonus = t.char.proficiency("intimidation") + 1
    will_dc = 10+t.creature["will"]
    roll = make_check("Demoralize", intimid_bonus, will_dc)
    frightened = roll.apply(lambda x: x if x >= Roll.Success else 0).expected()
    print("Frightened:", frightened)

    bon_mot = make_check("Bon mot", t.char.proficiency("diplomacy"), will_dc)
    penalty = bon_mot.apply(bon_mot_pen)
    print("Bon Mot Penalty:", penalty.expected())

    roll = make_check("Demoralize2", -penalty+intimid_bonus, will_dc)
    frightened = roll.apply(lambda x: x if x >= Roll.Success else 0).expected()
    print("Frightened:", frightened)

def weapon_damage():
    t = StatsTest()
    def calc_damage(level):
        t.set_level(level)
        # 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 t.char.weapons:
            attack = Attack(t.char, t.creature, weapon, sneak=False)
            res[weapon.name] = attack.strike().expected()
        # res["hp"] = creature["hp"]
        return res

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

def twin_takedown_test():
    t = StatsTest()
    # t.creature_plan["ac"] = Difficulty.Extreme
    def calc_damage(level):
        t.set_level(level)
        res = dict()
        wp1 = t.get_weapon("tricky pick")
        wp2 = t.get_weapon("gauntlet")

        # Twin takedown
        attack = Attack(t.char, t.creature, wp1, sneak=True)
        dmg = attack.strike().expected()
        attack.previous_attacks += 1
        attack.weapon = wp2
        dmg += attack.strike().expected()
        res["Twin Takedown"] = dmg

        # Double slice
        attack = Attack(t.char, t.creature, wp1, sneak=True)
        dmg = attack.strike().expected()
        attack.sneak = False
        attack.weapon = wp2
        dmg += attack.strike().expected()
        res["Double Slice"] = dmg

        return res

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

def multiple_attacks():
    t = StatsTest()
    t.set_level(5)
    wp = t.get_weapon("Shortsword")
    wp.agile = False
    # t.creature["ac"] -= 1
    attack = Attack(t.char, t.creature, wp, sneak=True)
    print(attack)
    for _ in range(4):
        attack.previous_attacks += 1
        print(attack)



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__':
    # weapon_damage()
    twin_takedown_test()
    # multiple_attacks()
    # demoralize_test()