Initial

2 files changed, 251 insertions, 0 deletions

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..6f2a68e
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,4 @@
+venv/
+.idea/
+__pycache__/
+*.py[codz]
diff --git a/main.py b/main.py
new file mode 100644
index 0000000..bf36ea8
--- /dev/null
+++ b/main.py
@@ -0,0 +1,247 @@
+import copy
+import math
+import random
+import numpy as np
+from collections import defaultdict
+import matplotlib.pyplot as plt
+import matplotlib.ticker as ticker
+from enum import Enum
+
+
+class Distribution:
+    def __init__(self, prob_table=None):
+        self.probabilities = prob_table
+
+    @property
+    def probabilities(self):
+        return self._prob_table
+
+    @probabilities.setter
+    def probabilities(self, value):
+        if value is None:
+            self._prob_table = defaultdict(lambda: 0.0)
+        elif isinstance(value, list) or isinstance(value, tuple):
+            self._prob_table = defaultdict(lambda: 0.0, {value[0], value[1]})
+        elif isinstance(value, dict) or isinstance(value, defaultdict):
+            self._prob_table = defaultdict(lambda: 0.0, value)
+        else:
+            raise ValueError("Unknown value type", type(value))
+
+    def expected(self):
+        ex = 0
+        for val, prob in self.probabilities.items():
+            ex += val * prob
+        return ex
+
+    def variance(self):
+        # Squared deviation from the mean
+        return (self + (-self.expected())).squared()
+
+    def standard_deviation(self):
+        return math.sqrt(self.variance().expected())
+
+    def squared(self):
+        return self.apply(lambda x: x*x)
+
+    @classmethod
+    def convert(cls, value):
+        if isinstance(value, cls):
+            return value
+        res = cls()
+        if isinstance(value, int) or isinstance(value, float):
+            res.probabilities = {value: 1.0}
+        else:
+            res.probabilities = value
+        return res
+
+    def prob_table(self):
+        return tuple(zip(*self._prob_table.items()))
+
+    def normalize(self):
+        summed = sum(self.probabilities.values())
+        self.probabilities = {k: v/summed for k, v in self.probabilities.items()}
+
+    def roll(self, count=1):
+        values, probs = self.prob_table()
+        res = random.choices(values, weights=probs, k=count)
+        return sum(res)
+
+    def show(self):
+        values, probs = self.prob_table()
+        self.__show_plot(values, probs, "Value", "Probability")
+
+    def test(self, sample_count):
+        values = list(self.probabilities.keys())
+        counts = defaultdict(lambda: 0)
+        for _ in range(sample_count):
+            counts[self.roll()] += 1
+
+        self.__show_plot(values, [counts[v] for v in values], "Value", "Times rolled")
+
+    @staticmethod
+    def __show_plot(xax, yax, xlb, ylb):
+        fig, ax = plt.subplots()
+        ax.bar(xax, yax)
+        ax.set_ylabel(ylb)
+        ax.set_xlabel(xlb)
+        ax.xaxis.set_major_locator(ticker.FixedLocator(xax))
+        # ax.yaxis.set_major_locator(ticker.FixedLocator(list(set(yax))))
+        plt.grid(True)
+        plt.show()
+
+    def insert(self, value, probability=1.0):
+        for val, prob in self.convert(value).probabilities.items():
+            self._prob_table[val] += probability * prob
+
+    def accumulate_prob(self, pred):
+        res = 0
+        for val, prob in self.probabilities.items():
+            if pred(val):
+                res += prob
+        return res
+
+    def apply(self, operator, *args, **kwargs):
+        res = Distribution()
+        for lv, lprob in self.probabilities.items():
+            res.insert(operator(lv, *args, **kwargs), lprob)
+        return res
+
+    def apply2(self, operator, other, *args, **kwargs):
+        res = Distribution()
+        for lv, lprob in self.probabilities.items():
+            for rv, rprob in self.convert(other).probabilities.items():
+                res.insert(operator(lv, rv, *args, **kwargs), lprob*rprob)
+        return res
+
+    def __add__(self, other):
+        return self.apply2(lambda x, y: x + y, other)
+
+    def __neg__(self):
+        return self.apply(lambda x: -x)
+
+    def __rmul__(self, count):
+        if not isinstance(count, int):
+            raise ValueError("Multiplication allowed only for int")
+        if count == 0:
+            return Distribution({0, 1.0})
+        res = copy.deepcopy(self)
+        if count < 0:
+            res = -res
+            count = -count
+        for _ in range(count-1):
+            res += self
+        return res
+
+
+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()