import copy
import math
import random
from collections import defaultdict
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker


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 __iadd__(self, other):
        self = self.apply2(lambda x, y: x + y, other)
        return self

    def __add__(self, other):
        return self.apply2(lambda x, y: x + y, other)

    __radd__ = __add__

    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