import numpy as np
from scipy import stats
from scipy import special as fn
from matplotlib import pyplot as plt
import collections, itertools
from operator import itemgetter as at


class BetaBinomial:
    __slots__ = ["T", "F"]
    def __init__(self, *args):
        if not any(args):
            #uninformative prior
            self.T = self.F = 1
        elif len(args)==1:
            #assuming rate
            self.T = args[0]*100.0
            self.F = (1-args[0])*100.0
        elif len(args)==2:
            self.T = args[0]
            self.F = args[1]
        else:
            raise SyntaxError("Illegal number of arguments")
    def update(self, *args):
        if len(args)==1:
            n = p = 0
            for x in args[0]:
                if x:
                    p += 1
                else:
                    n += 1
            return BetaBinomial(self.T + p, self.F + n)
        elif len(args)==2:
            return BetaBinomial(self.T + args[0], self.F + args[1])
        else:
            raise SyntaxError("Illegal number of arguments")
    def pdf(self, x):
        return stats.beta.pdf(x, self.T, self.F)
    def cdf(self, x):
        return stats.beta.cdf(x, self.T, self.F)
    def posterior(self, l, u):
        if l>u:
            return 0.0
        return self.cdf(u)-self.cdf(l)
    def mean(self, n=1):
        return self.T * n / (self.T + self.F)
    def plot(self, l=0.0, u=1.0):
        x = np.linspace(u, l, 1001)
        y = stats.beta.pdf(x,self.T, self.F)
        y=y/y.sum()
        plt.plot(x, y)
        plt.xlim((l,u))
    def predict(self, t, f, log=False):
        a = self.T
        b = self.F
        log_pmf = (fn.gammaln(t+f+1) + fn.gammaln(t+a) + fn.gammaln(f+b) + fn.gammaln(a+b)) - \
        (fn.gammaln(t+1) + fn.gammaln(f+1) + fn.gammaln(a) + fn.gammaln(b) + fn.gammaln(t+f+a+b))
        if log:
            return log_pmf
        return np.exp(log_pmf)

class BetaBernoulli(BetaBinomial):
    def update(self, *args):
        if len(args)==1:
            n = p = 0
            for x in args[0]:
                if x:
                    p += 1
                else:
                    n += 1
            return BetaBernoulli(self.T + p, self.F + n)
        elif len(args)==2:
            return BetaBernoulli(self.T + args[0], self.F + args[1])
        else:
            raise SyntaxError("Illegal number of arguments")
