import numpy as np
import scipy
import scipy.optimize

import mono_cubic2 as mc

def resize_with_const(v, n, val=0):
    """
    Resize a 1D vector to a specified length `n`.

    If the input vector `v` is longer than `n`, it is truncated.
    If it is shorter, it is padded with the constant value `val`.
    If it is already of length `n`, it is returned unchanged.

    Parameters:
        v (array-like): Input 1D vector (list or NumPy array).
        n (int): Target length of the output vector.
        val (scalar, optional): Value used to pad if `v` is shorter than `n`. Default is 0.

    Returns:
        np.ndarray: Resized 1D NumPy array of length `n`.
    """
    v = np.asarray(v)
    
    if len(v) == n: return v
    
    if len(v) > n: return v[:n]
    
    return np.concatenate([v, np.full(n - len(v), val)])

def safe_exp(x, max_exp = 700):
    """
    A numerically robust version of np.exp that avoids overflow by clipping the input.
    
    Parameters:
        x : array_like
            Input value or array.
        max_exp : float
            Maximum allowed exponent value. np.exp(709) ≈ 8.2e307 (close to float64 max).
    
    Returns:
        array_like
            The exponential of the input with overflow protection.
    """

    return np.exp(np.clip(x, -max_exp, max_exp))

def safe_expit(x, max_exp = 700): return 1/(1 + safe_exp(-x, max_exp))

"""
    Fitting data

        {(x_i, y_i) : x_i in R, y_i in {0,1}, i = 0, ..., n-1 }
    
    to model function 

        f(x|pars) = 1/(1 + exp(-F(x|beta)))         beta = beta(pars)

    with log odds of polynomial form:

        log(f(x|pars)/(1 - f(x|pars))) = F(x|beta)  beta = beta(pars)

    where F is decision function (aka logit)
        
        F(x|beta) = sum_{i=0}^degree beta_i x^i
                
    and coefficients
        
        beta = [beta_i(pars)]_{i=0}^degrees
    
    Conditional probability

        Prob(Y = y_i|x, pars) = 1/(1 + exp(-s(y) F(x|beta(pars))) 

    with

            s(y) = 2*y -1
"""
class LogisticPolyRegression:

    """
        Constructor

        Input:
            degree: int, degree of polynomial
            mono: boolean, default False
            lambda: None or tuple float, L1 and L2 regularization

    """
    def __init__(self, degree = 1, mono = False, lam = None):

        self.degree = degree
        self.mono = mono
        self.big = 200
        self.small = 1e-8
        self.lam = lam

        if mono:
            assert self.degree in [1, 3], f"Degree {self.degree} not supported in mono!"

        self.mono3 = self.mono and (self.degree == 3)

    """
        Mapping regression parameters pars to coefficients beta 
            
            beta = beta(pars)
        
        used in decision function:

            F(x|beta) = sum_{i=0}^degree beta_i x^i

        Input:
            pars
        
        Return:
            beta
    """
    def get_beta(self, pars):
        
        beta = mc.forward_map(pars) if self.mono3 else pars
        return np.array(beta)
    
    """
        Mapping beta to regression parameters used in decision function.

        Input:
            beta: coefficient beta
        
        Return:
            pars
    """
    def get_pars(self, beta):

        pars = mc.backward_map(beta) if self.mono3 else beta
        return np.array(pars)

    """
        Calculate jacobian between beta and  regression parameters
            
            J = d(beta)/d(pars)
              = [d(beta_i)/d(pars_a)]_{i,a}
        and
            
            H = [d^2 beta_i/(d(pars_a) d(pars_b))]_{i,a,b}

        Input:
            pars
            hess: boolean, False
        
        Return:
            J       if hess = True
            (J, H)  if hess = False
    """
    def get_jac_beta(self, pars, hess = False):
        
        n = len(pars)
            
        J = mc.forward_map_jacobian(pars) if self.mono3 else np.eye(n)
            
        if not hess: return J
        
        H = mc.forward_map_hessian(pars) if self.mono3 else np.zeros(shape = (n, n, n))

        return (J, H)


    """
        Model function 

            f(x|pars) = 1/(1 + exp(-F(x|beta)))  beta = beta(pars)

        Input:
            x: scalar value or a array of values
            pars: array of r = degree + 1 floats, model parameters array of floats

        Return:
            model function values 
    """
    def model(self, x, pars):

        beta = self.get_beta(pars)  # beta
        X = np.column_stack([x**i for i in range(len(beta))])
        F = X @ beta              # decision function, X beta

        return safe_expit(F)

    """
        Calculate negative log-likelihood function

            nllf = -sum_i log(Prob(Y = y_i| x_i, pars))   : neg. log likelihood
            
            grad = [d(nllf)/d(pars_a)]_a                  : jacobian
        
        where

            Prob(Y = y_i| x_i, pars) = 1/(1 + exp(-s_i F(x_i| beta)))  beta=beta(pars)
            s_i = 2*y_i  - 1
     
        Input:
            x: array of n floats
            y: array of n int in {0,1}  
            pars: array of r = degree + 1 floats, model parameters array of floats
            jac: boolean, default False, if jacobian is needed
        
        Return:
            nllf           : if jac is false
            (nllf, grad)   : if jac is true
    """
    def nllf(self, x, y, pars, jac = False):

        beta = self.get_beta(pars)
        X = np.column_stack([x**i for i in range(len(beta))])
    
        # signs
        s = 2.0*y - 1
        
        # decision function for conditional probability Prob(Y = y| x)
        F = s*(X @ beta)

        nllf = np.sum(np.log(1 + safe_exp(-F))) 

        if self.lam is not None:
            nllf += self.lam[0]*np.sum(np.abs(pars)) + self.lam[1]*np.sum(pars**2)

        if not jac: return nllf
        
        J = self.get_jac_beta(pars)

        grad = -(s*safe_expit(-F)) @ (X @ J)

        if self.lam is not None:
            grad += self.lam[0]*np.sign(pars) + 2*self.lam[1]*pars

        return (nllf, grad)
    
    """
        Estimate parameters.

        Input:
            x: array of n floats
            y: array of n int in {0,1}
        
        Return:
            pars0
    """
    def get_est_pars(self, x, y):

        L = np.log(2*len(x) + 1)

        if self.mono3:    
            z = L*np.polyfit(x, 2.0*y - 1, 1)[::-1]
            beta = resize_with_const(z, self.degree + 1, 1e-8)    
        else:
            beta = L*np.polyfit(x, 2.0*y - 1, self.degree)[::-1]

        return self.get_pars(beta)
    
    """
        Performing logistic regression with log odds of polynomial form:

            log(f(x|pars)/(1 - f(x|pars))) = F(x|beta)  beta = beta(pars)
        
        and this gives           
            
            f(x|pars) = 1/(1 + exp(-F(x|beta)))

        where coefficient beta = [beta_i(pars)]_{i=0}^degree, with decision 
        function (aka logit)

            F(x|beta) = sum_{i=0}^degree beta_i x^i
                
        Input:
            x: array of n floats
            y: array of n int in {0, 1}
        
        Return:
            pars
    """
    def fit(self, x, y, pars0 = None, method = "local"):
        
        bnds = [(-self.big, self.big)]*(self.degree + 1)
  
        if method == "local":

            pars0 = self.get_est_pars(x, y)
            cost = lambda pars: self.nllf(x, y, pars, jac = True)
            res = scipy.optimize.minimize(cost, x0 = pars0, method = 'L-BFGS-B', jac = True, bounds = bnds, tol=1e-12)

        elif method == "diff_evol":
            
            cost = lambda pars: self.nllf(x, y, pars, jac = False)
            res = scipy.optimize.differential_evolution(cost, bounds = bnds, tol = 1e-8, polish=False)

            cost = lambda pars: self.nllf(x, y, pars, jac = True)
            res = scipy.optimize.minimize(cost, x0 = res.x, method = 'L-BFGS-B', jac = True, bounds = bnds, tol=1e-12)

        elif method == "anneal":

            cost = lambda pars: self.nllf(x, y, pars, jac = False)
            res = scipy.optimize.dual_annealing(cost, bounds = bnds)
        
        else:
            assert False, "This method is not supported."
        
        return {"pars": res.x, "cost": res.fun, "success": res.success}

    """
        Producing goodness of fit measures:
        
            LLF = log_likelihood function
            AIC = Akaike information criterion
            BIC = Bayesian information criterion

        Input:
            x: array of n floats
            y: array of n int in {0,1}
            pars: array of r = degree+1 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)}
        
        Ref:
            https://en.wikipedia.org/wiki/Logistic_regression
            https://en.wikipedia.org/wiki/Akaike_information_criterion
            https://www.medicine.mcgill.ca/epidemiology/joseph/courses/epib-621/logfit.pdf
    """
    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) + self.small)
        chi2 = np.sum(r**2)
        p_val = scipy.stats.chi2.sf(chi2, dof)

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

        # accuracy A
        A = np.count_nonzero(matches)/n

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

    """
        Calculation of asymptotic variance-covariance matrix of parameters pars

           cov_{asymp}[pars] = H^{-1}

        where H is hessian of nllf 

           H = [d^2(nllf)/(d(pars)_a d(pars)_b ]_{a,b}
        
        for the logistic regression of the polynomial model:
        
            log(f(x)/(1 - f(x))) ~ sum_{i=0}^degree b_i(pars) x^i

        Input:
            x: array of n floats
            pars: array of r = degree+1 floats, model parameters
            
        Return:
            array of rxr floats; r = degree + 1
        
        Ref:
            https://stats.stackexchange.com/questions/89484/how-to-compute-the-standard-errors-of-a-logistic-regressions-coefficients
            https://goodboychan.github.io/machine_learning/2020/09/14/02-Regularized-likelihood-methods.html
    """
    def cov(self, x, y, pars):
        
        # coefficients
        beta = self.get_beta(pars)
        
        # design matrix -- add column of 1's at the beginning of your X_train matrix
        X = np.column_stack([x**i for i in range(len(beta))])

        # Jacobian J = [dbeta_i/dpars_j]_{ij}
        J, H = self.get_jac_beta(pars, hess = True)

        # signs
        s = 2.0*y - 1
        
        # decision function for conditional probability Prob(Y = y| x)
        F = s*(X @ beta)

        # probabilities p_i = P(Y=y_i | x_i)
        p = safe_expit(F)
        q = 1 - p

        # calculate hessian
        L = X @ J
        H = (L.T*(q*p))@L - np.tensordot((s*q)@X, H, axes = ([0], [0]))

        if self.lam is not None: 
            Hp = H + 2*self.lam[1]*np.eye(len(pars))   # H' = H + lambda id 
            iHp = np.linalg.inv(Hp)                    # inv(H')

            return iHp@H@iHp
        
        # covariance matrix C_params = H^-1
        return np.linalg.inv(H)


    """
        Calculating quantiles of the model parameters at given probabilities p 
        for normal distribution of parameters:

            pars ~ N(mean_pars, cov_pars)
        
        Input:
            probs: array of m floats, probabilities
            mean_pars: array of r = degree+1 floats, mean model parameters
            cov_pars: array of rxr floats, variance-covariance matrix of parameters
        
        Return:
            array of mxr floats
    """
    def get_pars_quantiles_normal(self, probs, mean_pars, cov_pars):
        
        # mean and standard variance parameters
        locs = mean_pars
        scales = np.sqrt(np.diag(cov_pars))

        # computing quantiles of parameters
        return locs + np.outer(scipy.stats.norm.ppf(probs), scales)

    """
        Calculating quantiles of the model values 
        
            f(x|pars) = 1/(1 + exp(-F(x|beta)))     beta = beta(pars)
        
        with

            F(x|beta) = sum_{i=0}^degree x^i beta_i
            
        at given probabilities p and values x assuming 
        normal distribution of parameters:

            pars ~ N(mean_pars, cov_pars)
        
        This distribution is asymptotic MLE distribution of parameters.

        Input:
            x: array of n float
            probs: array of m floats, probabilities
            mean_pars: array of r = degree+1 floats, mean model parameters
            cov_pars: array of rxr floats, variance-covariance matrix of parameters
        
        Return:
            array of mxn floats
    """

    def get_model_quantiles_normal(self, x, probs, mean_pars, cov_pars, 
                                   exact = True, seed = 1977, m = 10**5):
        
        mean_beta = self.get_beta(mean_pars)
        X = np.column_stack([x**i for i in range(len(mean_beta))])
        
        if exact and self.mono3:

            # init random generator
            rng = np.random.default_rng(seed)
            pars = rng.multivariate_normal(mean_pars, cov_pars, size = m)

            # get betas
            beta = np.apply_along_axis(self.get_beta, 1, pars)
            
            # quantiles of decision function
            Q = np.quantile(X@beta.T, probs, axis = 1)
            
            return np.apply_along_axis(safe_expit, 1, Q)

        # J = d(beta)/d(pars)
        J = self.get_jac_beta(mean_pars)
        
        # transform data
        S = X@J

        # mean and standard variance of logit (aka log of odds)
        locs = X@mean_beta
        scales = np.sqrt(np.diag(S@cov_pars@S.T))

        # computing quantiles of logit
        Q = locs + np.outer(scipy.stats.norm.ppf(probs), scales)
        
        # convert logit to expit
        return safe_expit(Q)

    """
        Calculating quantiles using delta method of the model values 
        
            f(x|pars) = 1/(1 + exp(-F(x|beta))) beta = beta(pars)

        with

            F(x|beta) = sum_{i=0}^degree x^i beta_i

        at given probabilities p and values x assuming 
        normal distribution of parameters :

            pars ~ N(mean_pars, cov_pars)

        This distribution is asymptotic MLE distribution of parameters.
        We approximate exact model with linear expansion

            f(x|pars) = p(x|pars_mean) + dp/db(x|pars_mean) (pars - pars_mean)
        
        and the last term is normally distributed.

        Input:
            x: array of n float
            probs: array of m floats, probabilities
            mean_pars: array of r = degree+1 floats, mean model parameters
            cov_pars: array of rxr floats, variance-covariance matrix of parameters 
        
        Return:
            array of mxn floats
    """
    def get_model_quantiles_delta(self, x, probs, mean_pars, cov_pars):
        
        mean_beta = self.get_beta(mean_pars)
        X = np.column_stack([x**i for i in range(len(mean_beta))])
        F = X@mean_beta

        # J = d(beta)/d(pars)
        J = self.get_jac_beta(mean_pars)

        # S = d(F)/d(pars)
        S = X@J

        # attributes of normal distribution of model values
        locs = safe_expit(F)
        scales = np.sqrt(np.diag(S@cov_pars@S.T))/(4*np.cosh(F/2)**2)

        # computing quantiles of logit
        Q = locs + np.outer(scipy.stats.norm.ppf(probs), scales)

        return np.clip(Q, a_min = 0, a_max = 1)

    """
        Generate parameters assuming normal distribution.

        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
    
        Return:
            array of mx(degree + 1)
    """
    def get_normal_pars(self, x, y, m, seed = 1977):
        
        res = self.fit(x, y, method = "diff_evol")

        assert res["success"], "Fit did not succeed."

        # optimal parameters
        pars = res["pars"]

        # covariance matrix of parameters
        cov = self.cov(x, y, pars)

        # init random generator
        rng = np.random.default_rng(seed)

        return rng.multivariate_normal(pars, cov, size = m)

    """
        Generate m parameters via non-parametric bootstrapping with a minimal constraint 
        that both groups should be present in the sampled data:

        boostrapped sample = (xb, yb)  by sampling with replacement pairs (x_i, y_i) 
                                    with condition that yb can not be just 0 or just 1 

        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
    
        Return:
            array of mx(degree + 1)
    """
    def get_nonparam_boots_pars(self, x, y, m, seed = 1977):
        
        # fitting original data
        res_fit = self.fit(x, y, method = "diff_evol")
        assert res_fit["success"]

        pars0 = res_fit["pars"]

        n = len(x)

        rng = np.random.default_rng(seed)

        # generate parameters
        lst = [pars0]
        while True:

            # create set indices for sampling with replacement + constraint
            idx = rng.choice(n, n)
            if np.sum(y[idx]) in [0, n]: continue

            # do fitting 
            res_fit = self.fit(x[idx], y[idx], pars0, method="local")
            if not res_fit["success"]: continue
            
            # store pars
            lst.append(res_fit["pars"])
            if len(lst) == m: break

        return np.array(lst)


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

            boostrapped sample = (xb, yb)  
                
                xb = (sampled with replacement from x0, sampled with replacement from x1)
                yb = (0...0, 1...1)

        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
        
        Return:
            array of mx(degree + 1)
    """
    def get_nonparam_stratified_boots_pars(self, x, y, m, seed = 1977):

        # pars of original data
        res_fit = self.fit(x, y, method = "diff_evol")
        assert res_fit["success"]

        pars0 = res_fit["pars"]
        
        # statistics about 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)])

        rng = np.random.default_rng(seed)

        # generate parameters
        lst = [pars0]
        for _ in range(m):
            # stratified sampling with replacement
            xb = np.concatenate([rng.choice(xs[i], ns[i]) for i in range(2)])

            # do fitting 
            res_fit = self.fit(xb, yb, pars0, method="local")
            if not res_fit["success"]: continue
            
            # store pars
            lst.append(res_fit["pars"])
            if len(lst) == m: break

        return np.array(lst)

    """
        Generate m parameters via parametric bootstrapping:

            boostrapped sample = (x, yb) yb ~ B(model(x, fitted pars))  
        
        where B is Bernoulli distribution

        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
        
        Return:
            array of mx(degree + 1)

        Ref:

        * https://www.scirp.org/journal/paperinformation?paperid=70962
        * https://en.wikipedia.org/wiki/Bernoulli_distribution
    """
    def get_parametric_boots_pars(self, x, y, m, seed = 1977):

        # first discuss original dataset
        res_fit = self.fit(x, y, method = "diff_evol")
        assert res_fit["success"]

        pars0 = res_fit["pars"]

        p = self.model(x, pars0)
    
        n = len(x)
        
        rng = np.random.default_rng(seed)

        # generate parameters
        lst = [pars0]
        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
            
            # do fitting 
            res_fit = self.fit(x, y_sim, pars0, method="local")
            if not res_fit["success"]: continue
            
            # store pars
            lst.append(res_fit["pars"])
            if len(lst) == m: break

        return np.array(lst)