import numpy as np
import scipy

"""
    Setting up scipy distribution to be used in Bayesian model

    Input:
        distr_str: string of scipy distribution
        n_pars: int, number of parameters
        parse_pars: function(pars) -> dict
    
    Return: 
        (log_pdf, distr_sample)
"""
def setup_scipy_distr(distr_str, n_pars, parse_pars):
    

    # log of pdf for all choice
    def log_pdfs(x, pars, choice = 2):
        logpdf = eval(distr_str).logpdf
        if choice in [0,1]: return logpdf(x, **parse_pars(pars))
        return (logpdf(x, **parse_pars(pars)), logpdf(x, **parse_pars(pars[n_pars:])))

    # sample w.r.t. distribution
    def distr_sample(rng, pars, n, choice = 2):    
        rvs = eval(distr_str).rvs
        if choice in [0,1]: return rvs( **parse_pars(pars), size = n, random_state = rng)
        return np.concatenate((rvs(**parse_pars(pars), size = n[0], random_state = rng), 
                            rvs(**parse_pars(pars[n_pars:]), size = n[1], random_state = rng)))
    
    return log_pdfs, distr_sample

"""
    Bayesian model describing conditional probability 

        Prob(X|Y = 1) 
            = Prob(Y=1)p(X|Y=1)/(Prob(Y=0) p(X|Y=0) + Prob(Y=1) p(X|Y=1))
            = 1 /(1 + O p(X|Y=0)/p(X|Y=1)) 
            = 1/(1 + exp( -F))

    where F is the decision function

        F =  log(p(X|Y=1)) - log(p(X|Y=0)) - log(O)
    
    and O are the odds

        O = Prob(Y=0)/Prob(Y=1)
"""
class BayesianModelRegression:
    
    """
        Constructor

        Input:
            odds: float, ratio Prob(Y=0)/Prob(Y=1)
            log_pdfs: function (x, pars, choice = 2) 
                        match choice:
                            case 0: return log_pdf0
                            case 1: return log_pdf1
                            case _: return (log_pdf0, log_pdf1)
            bounds: tuple of bounds, (bounds0, bound1)

            distr_sample: function (rng, x, pars, n, choice):
                          generate n sampled os points using pdfs(choice, pars) 
    """   
    def __init__(self, odds, log_pdfs, bounds, distr_sample = None):
        
        self.odds = odds
        self.log_odds = np.log(odds)
        self.log_pdfs = log_pdfs
        self.bounds = bounds
        self.distr_sample = distr_sample

    """
        Calculate decision function:

             decision = log(p(X|Y=1)) - log(p(X|Y=0)) - log(odds)

        Input:
            x: float or array of floats
            pars: parameters for log_pdfs
        Return:
            float or array of float
    """
    def decision(self, x, pars):

        # log of pdf for each group
        lf0, lf1 = self.log_pdfs(x, pars)
        
        return lf1 - lf0 - self.log_odds
    
    """
        Calculate model of the conditional probability Prob(X|Y = 1)
        Input:
            x: float or array of floats
            pars: parameters for log_pdfs
    """    
    def model(self, x, pars):
    
        # decision function
        F = self.decision(x, pars)

        # calculating model
        return 1/(1 + np.exp(-F))
    
    """
        Negative Log Likelihood function:

            neg. log likelihood = -sum_i log(Prob(X = x_i, Y = y_i))
        
        where

            Prob(X, Y = 1) = 1/(1 + exp(-F))
            Prob(X, Y = 0) = 1 -  Prob(X, Y = 1) = 1/(1 + exp(+F))
        
        with

            log Prob(X, Y = y) = -log(1 + exp(-S(y) F))
            S(y) =  [ +1 : y = 1
                    [ -1 : y = 0

        Input:
            x: array of floats
            y: array of ints in {0,1}
            pars: array of floats, model parameters
        
        Return:
            float: negative log likelihood

    """
    def nllf(self, x, y, pars):
        
        # signs for the groups:
        # group 1 has + sign and group 0 has - sign
        S = 2.0*y - 1

        # decision function
        F = self.decision(x, pars)
        
        return np.sum(np.log(1 + np.exp(-S*F)))
    

    """
        MLE fitting of a distribution,  given by log_pdf, to data x associated 
        to the group 0 or 1 by maximizing 

            loglikehood_{single group} = sum_i log_pdf(x | pars)

        Input:
            x: array of floats
            choice: int in {0,1}, selecting the group
            method: string in ["local", "diff_evol", "anneal"]
            seed: int, seed of the random generator
        
        Return:
            {"pars": pars_MLE, "cost": NLLF at pars_MLE}
    """
    
    def fit_distr(self, x, choice, method = "local", seed = 1977):
        fname = "fit_distr"

        cost = lambda pars: -np.sum(self.log_pdfs(x, pars, choice))
        bnds = self.bounds[choice]

        match method:
            case "local":
                # random parameters from boundaries
                pars0 = np.random.default_rng(seed).uniform(*zip(*bnds))
                # use local optimizer
                res = scipy.optimize.minimize(cost, pars0, bounds = bnds, method = "L-BFGS-B")
            case "diff_evol":
                res = scipy.optimize.differential_evolution(cost, bounds = bnds)
            case "annel":
                res = scipy.optimize.dual_annealing(cost, bounds = bnds)
            case _:
                assert False, f"{fname}::this method does not exist"
    
        return {"pars": res.x, "success": res.success, "cost": res.fun} 


    """
        MLE fitting of the Bayesian model:

            pars_MLE = argmin_pars NLLF(pars| x, y)

        Input:
            x: array of floats
            y: array of ints in {0,1}
            method: string in ["local", "diff_evol", "anneal"], optimizer
        
        Return:
            {"pars": pars_MLE, "cost": NLLF at pars_MLE}
    """
    def fit(self, x, y, pars0 = None, method = "local"):
        fname = "fit"

        # defined nllf as function of parameters, data is already included
        cost = lambda pars: self.nllf(x, y, pars)

        # joint bounds of two groups
        bnds = np.concatenate(self.bounds)
        
        match method:
            case "local":
                # estimate initial guess of parameters (for local method)
                if pars0 is None:
                    get_pars = lambda choice: self.fit_distr(x[y == choice], choice)["pars"]
                    pars0 = np.r_[get_pars(0), get_pars(1)]
          
                # optimize using local optimizer
                res = scipy.optimize.minimize(cost, pars0, bounds = bnds, method = "L-BFGS-B")
            case "diff_evol":
                res = scipy.optimize.differential_evolution(cost, bounds = bnds)
            case "anneal":
                res = scipy.optimize.dual_annealing(cost, bounds = bnds)
            case _:
                assert False, f"{fname}::this method does not exist"

        return {"pars": res.x, "success": res.success, "cost": res.fun} 

    """
        Producing goodness of fit measures:
    
            LLF = log_likelihood function
            AIC = Akaike information criterion
            BIC = Bayesian information criterion
            
        Input:
            x: array of floats
            y: array of ints in {0,1}
            pars: array of floats,  model parameters
            thresh: float, default 0.5, threshold value for classification

        Return:
            {"n": n, "k":k, "dof":n-k, 
            "LLF": log_likelihood, 
            "AIC": AIC, 
            "BIC": BIC, 
            "A": classification accuracy (threshold values = 0.5 prob)}    
    """
    def goodness_of_fit(self, x, y, pars, thresh = 0.5):

        # model probabilities
        p = self.model(x, pars)

        # log likelihood
        llf = -self.nllf(x, y, pars)
        
        # information criteria
        k, n = len(pars), len(x)
        AIC = 2*k - 2*llf
        BIC = k*np.log(n) - 2*llf

        # chi2
        dof = n - k
        r = (y - p)/np.sqrt(p*(1-p))
        chi2 = np.sum(r**2)
        p_val = scipy.stats.chi2.sf(chi2, dof)

        # using model as classifier
        matches = y == np.heaviside(p - thresh, 1)

        return {"LLF": llf, "AIC": AIC, "BIC": BIC, 
                "A" : np.count_nonzero (matches)/n,
                "chi2": chi2, "p-value(chi2)": p_val,  # not very useful
                "n": n, "k": k, "dof": dof}

    """
        Generate m parameters via non-parametric bootstrapping with minimal constraint

            bootstrapped sampled = (xb, yb) sampled with replacement from (x,y)

        Input:
            x : array of n floats
            y : array of n int in {0,1}
            m: integer, number of samples
            seed : int, seed for the random generator
            method: string in ["local", "diff_evol", "anneal"], optimizer
        
        Return:
            array of m x len(pars) floats
    """
    def get_nonparam_boots_pars(self, x, y, m, seed = 1, method = "local"):
        fname = "get_nonparam_boots_pars"

        rng = np.random.default_rng(seed)

        # discussing original data 
        res = self.fit(x, y, method = method)
        assert res["success"], f"{fname}::fitting original data failed."

        # generate bootstrapped parameters
        pars = res["pars"] 
        lst = [pars]
        n = len(x)
        
        while True:
            
            # sampling with replacement with restrictions
            idx = rng.choice(n, n)
            if np.sum(y[idx]) in [0, n]: continue
            
            res = self.fit(x[idx], y[idx], pars0 = pars, method = method)
            if not res["success"]: continue

            lst.append(res["pars"])
            if len(lst) == m: break

        return np.array(lst)


    """
        Generate m parameters via non-parametric stratified bootstrapping:

            bootstrapped sampled = (xb, yb) sampled with replacement from (x,y) for each groups separately

        meaning         
          
            xb = (sampled with replacement from x0, sampled with replacement from x1)
            yb = (0 ... 0, 1 ... 1)
                    n0        n1
        
        Note samples from each group in yb is constant and same as in y.

        Input:
            x : array of n floats
            y : array of n int in {0,1}
            m: integer, number of samples
            seed : int, seed for the random generator
            method: string in ["local", "diff_evol", "anneal"], optimizer
        
        Return:
            array of m x len(pars) floats
    """
    def get_nonparam_strat_boots_pars(self, x, y, m, seed = 1, method = "local"):
        fname = "get_nonparam_strat_boots_pars"

        rng = np.random.default_rng(seed)

        # discussing original data 
        res = self.fit(x, y, method = method)
        assert res["success"], f"{fname}::fitting original data failed."

        # separate data of both groups
        xs = [x[y == i] for i in range(2)]
        ns = [len(e) for e in xs]
        
        # common vector states
        yb = np.concatenate([np.full(ns[i], i) for i in range(2)])
        
        # generate bootstrapped parameters
        pars = res["pars"]   
        lst = [pars]
        while True:
            # stratified sampling with replacement
            xb = np.concatenate([rng.choice(xs[i], ns[i]) for i in range(2)])
        
            # do fitting
            res = self.fit(xb, yb, pars0 = pars, method = method)
            if not res["success"]: continue

            lst.append(res["pars"])
            if len(lst) == m: break

        return np.array(lst)
    
    """
        Generate m parameters via parametric bootstrapping:
            
            bootstrapped sample = (x, yb)  yb ~ B(ymodel)
        
        Input:
            x : array of n floats
            y : array of n int in {0,1}
            m: integer, number of samples
            seed : int, seed for the random generator
            method: 
        
        Return:
            array of m x len(pars) floats
    """
    def get_param_boots_pars(self, x, y, m, seed = 1, method = "local"):
        fname = "get_param_boots_pars"

        rng = np.random.default_rng(seed)
        
        # discussing original data 
        res = self.fit(x, y, method = method)
        assert res["success"], f"{fname}::fitting original data failed."

        # calculate predicted conditional probabilities 
        pars = res["pars"]
        p = self.model(x, pars)

        # generate bootstrapped parameters
        lst = [pars]
        n = len(x)

        while True:

            # Generate new binary outcomes from Bernoulli(p_i)
            y_sim = rng.binomial(n = 1, p = p)

            if np.sum(y_sim) in [0, n]: continue
            
            # fit and get new parameter
            res = self.fit(x, y_sim, pars0 = pars, method=method)
            if not res["success"]: continue

            lst.append(res["pars"])
            if len(lst) == m: break

        return np.array(lst)
    
    """
        Generate m parameters via parametric stratified bootstrapping by 
        sampling x from parametrized distributions associated to individual groups:

            bootstrapped sample = (xb, yb)
                xb = (sampled from distr for x0, sampled from distr for x1)
                yb = (0 ... 0, 1 ... 1)
                        n0        n1

        Input:
            x : array of n floats
            y : array of n int in {0,1}
            m: integer, number of samples
            seed : int, seed for the random generator
            method: 
        
        Return:
            array of m x len(pars) floats
    """
    def get_param_strat_boots_pars(self, x, y, m, seed = 1, method = "local"):
        fname = "get_param_strat_boots_pars"

        assert self.distr_sample is not None, f"{fname}::distr_sample is not defined"

        rng = np.random.default_rng(seed)
        
        # separate data of both groups
        xs = [x[y == i] for i in range(2)]
        ns = [len(e) for e in xs]
        
        # separate data of both groups
        pars_g = np.concatenate([self.fit_distr(e, i)["pars"] for i, e in enumerate(xs)])
        
        # discussing original data 
        res = self.fit(x, y, method = method)
        assert res["success"], f"{fname}::fitting original data failed."

        # common vector states
        yb = np.concatenate([np.full(ns[i], i) for i in range(2)])

        # generate bootstrapped parameters
        pars = res["pars"]
        lst = [pars]
        
        while True:

            # Generate new sample of points for each group
            xb = self.distr_sample(rng, pars_g, ns)
            
            # fit and get new parameter
            res = self.fit(xb, yb, pars0 = pars, method = method)
            if not res["success"]: continue

            lst.append(res["pars"])
            if len(lst) == m: break

        return np.array(lst)