import numpy as np
import scipy
import scipy.stats

"""
    Model function 

        p(x|b) = 1/(1 + exp(-F(x|b)))
    
    with log odds of polynomial form:

        log(p(x|b)/(1 - p(x|b))) = F(x|b)

        F(x|b) = sum_{i=0}^degree b_i x^i
                
    with decision function (aka logit) F and parameters 
        
        b = [b_i]_{i=0}^degree

    Input:
        x: scalar value or a array of values
        b = [b_i]_{i=0}^degree: array of r = degree + 1 floats
           model parameters array of floats

    Return:
        model function values 

"""
def logit_poly_model(x, b):

    X = np.column_stack([x**i for i in range(len(b))])
    F = X @ b

    return 1/(1 + np.exp(-F))


"""
    Performing logistic regression with log odds of polynomial form:

        log(p(x|b)/(1 - p(x|b))) = F(x|b)
    
        F(x|b) = sum_{i=0}^degree b_i x^i
                
    with decision function (aka logit) F and b = [b_i]_{i=0}^degree. 
    This gives 

        p(x|b) = 1/(1 + exp(-F(x|b)))
    
    Input:
        lm: instance linear_model.LogisticRegression
        x: array of n floats
        y: array of n floats
        d: degree of decision function
    
    Return:
        params = [b_0, ..., b_degree], array of r = degree + 1 floats
"""

def logit_poly_fit(lm, x, y, degree = 1):
   
   X_feature = np.column_stack([x**i for i in range(1, degree+1)])

   lm.fit(X_feature, y)

   return np.r_[lm.intercept_[0], lm.coef_[0,:]]

"""
    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}
        b: array of r = degree+1 floats, model parameters
    
    Return:
        {"n": n, "k":k, "dof":n-k, "LLF": log_likelihood, "AIC": AIC, "BIC": BIC}
    
    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 logit_poly_goodness_of_fit(x, y, b):
    
    # model probabilities
    p = logit_poly_model(x, b)

    # log likelihood
    eps = 1e-20  # prevent log(0)
    llf = np.sum(y*np.log(p + eps) + (1 - y)*np.log(1 - p + eps))
    
    # information criteria
    k, n = len(b), 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)

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

"""
    Calculation of asymptotic variance-covariance matrix for the 
    logistic regression of the polynomial model:
    
        log(p(x)/(1 - p(x))) ~ sum_{i=0}^degree b_i x^i

    Input:
        x: array of n floats
        b: 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
"""
def logit_poly_cov(x, b):
    
    # Calculate matrix of predicted class probabilities.
    probs = logit_poly_model(x, b)
    
    # 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(b))])

    # Initiate matrix of 0's, fill diagonal with each predicted observation's variance
    V = np.diagflat(probs*(1 - probs)) # dig.matrix where each element is p*(1-p)

    # Covariance matrix C_params = (X^T V X)^-1
    return np.linalg.inv(X.T@V@X)


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

        b ~ N(mean_b, cov)
    
    Input:
        probs: array of m floats, probabilities
        mean_b: array of r = degree+1 floats, mean model parameters
        cov: array of rxr floats, variance-covariance matrix of parameters
    
    Return:
        array of mxn floats
"""

def logit_poly_pars_quantiles_normal(probs, mean_b, cov):
    
    # mean and standard variance parameters
    locs = mean_b
    scales = np.sqrt(np.diag(cov))

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

"""
    Calculating quantiles of the model values 
    
        p(x|b) = 1/(1 + exp(-F(x|b)))
    
    with

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

        b ~ N(mean_b, cov)
    
    This distribution is asymptotic MLE distribution of parameters.

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

def logit_poly_model_quantiles_normal(x, probs, mean_b, cov):
    
    X = np.column_stack([x**i for i in range(len(mean_b))])

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

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

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

    with

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

        b ~ N(mean_b, cov)

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

        p(x|b) = p(x|b_mean) + dp/db(x| b_mean) (b - b_mean)
    
    and the last term is normally distributed.

    Input:
        x: array of n float
        probs: array of m floats, probabilities
        mean_b: array of r = degree+1 floats, mean model parameters
        cov: array of rxr floats, variance-covariance matrix of parameters 
    
    Return:
        array of mxn floats
"""
def logit_poly_model_quantiles_delta(x, probs, mean_b, cov):

    X = np.column_stack([x**i for i in range(len(mean_b))])
    F = X@mean_b

    locs = 1/(1 + np.exp(-F))
    scales = np.sqrt(np.diag(X@cov@X.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 m parameters via non-parametric bootstrapping with a minimal constraint 
    that both groups should be present in the sampled data.

     Input:
        lm: linear_model.LogisticRegression
        x : array of n floats
        y : array of n int in {0,1}
        m: integer, number of samples
        degree: int,  degree of decision function
        seed : int, seed for the random generator
    
    Return:
        array of mx(degree + 1)
    
    Return:
        array of mx(degree + 1)
"""
def get_nonparam_boots_pars(lm, x, y, m, degree = 1, seed = 1977):

    rng = np.random.default_rng(seed)
    
    # fitting original data
    pars = logit_poly_fit(lm, x, y, degree = degree)
    
    n = len(x)

    # generate parameters
    lst = [pars]
    while True:

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

        lst.append(logit_poly_fit(lm, x[idx], y[idx], degree=degree))
        if len(lst) == m: break

    return np.array(lst)


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

    Input: 
        lm: linear_model.LogisticRegression
        x : array of n floats
        y : array of n int in {0,1}
        m: integer, number of samples
        degree: int,  degree of decision function
        seed : int, seed for the random generator
    
    Return:
        array of mx(degree + 1)
"""
def get_nonparam_stratified_boots_pars(lm, x, y, m, degree = 1, seed = 1977):

    rng = np.random.default_rng(seed)

    # pars of original data
    pars = logit_poly_fit(lm, x, y, degree = degree)
    
    # 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)])

    # generate parameters
    lst = [pars]
    for _ in range(m):
        # stratified sampling with replacement
        xb = np.concatenate([rng.choice(xs[i], ns[i]) for i in range(2)])
        # do fitting
        lst.append(logit_poly_fit(lm, xb, yb, degree = degree))

    return np.array(lst)

"""
    Generate m parameters via parametric bootstrapping.

    Input: 
        lm: linear_model.LogisticRegression
        x : array of n floats
        y : array of n int in {0,1}
        m: integer, number of samples
        degree: int,  degree of decision function
        seed : int, seed for the random generator

    Return:
        array of mx(degree + 1)

    Ref:

      * https://www.scirp.org/journal/paperinformation?paperid=70962
"""
def get_parametric_boots_pars(lm, x, y, m, degree = 1, seed = 1977):

    # first discuss original dataset
    pars = logit_poly_fit(lm, x, y, degree = degree)
    p = logit_poly_model(x, pars)

    rng = np.random.default_rng(seed)

    n = len(x)

    # generate parameters
    lst = [pars]
    while True:

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

        if np.sum(y_sim) in [0, n]: continue
        
        lst.append(logit_poly_fit(lm, x, y_sim, degree = degree))

        if len(lst) == m: break

    return np.array(lst)