import numpy as np
import pandas as pd
import statsmodels.api as sm

try:
    from scikits.bootstrap import ci
except ModuleNotFoundError:
    print('Could not load scikits-bootstrap: bootstrap CIs not available.')

__all__ = ['kreg_perm',
            'argrank',
            'dist2kernel',
            'kernel2dist',
            'computePWDist',
            'computeKregStat']
            
"""Tested one kernel, logistic regression.
Multiple kernels and continuous outcome should be tested."""

def kreg_perm(y, Ks, X=None, binary=True, nperms=9999, seed=110820, returnPerms=False):
    """Kernel regression of Y adjusting for covariates in X
    Implemented the permutation method from:
    
    Zhao N, Chen J, Carroll IM, Ringel-Kulka T, Epstein MP, Zhou H, et al.
        Testing in microbiome-profiling studies with MiRKAT, the microbiome
        regression-based kernel association test. Am J Hum Genet. The
        American Society of Human Genetics; 2015;96(5):797-807.

    Parameters
    ----------
    y : pd.Series or np.ndarray shape (n,1)
        Endogenous/outcome variable.
        Can be continuous or binary (use binary=True)
    Ks : pd.DataFrame or np.ndaray shape (n,n) OR list
        Kernel or list of kernels
        (square, symetric and positive semi-definite)
    X : pd.DataFrame OR np.ndarray shape (n,i)
        Matrix of covariates for adjustment.
    binary : bool
        Use True for logistic regression.
    nperms : int
        Number of permutations for the test.
    returnPerms : bool
        If True and Ks is a single kernel,
        return the observed and permuted Q-statistics.

    Returns
    -------
    pvalue : float
        Permutation P-value or global/omnibus P-value,
        if there are multiple kernels
    pIndivid : ndarray shape (k)
        Optionally returns vector of p-values, one
        for each kernel provided."""

    def computeQ(K, resid, s2):
        return np.linalg.multi_dot((resid/s2, K, resid))

    n = len(y)
    if binary:
        family = sm.families.Binomial()
    else:
        family = sm.families.Gaussian()

    if X is None:
        X = np.ones(y.shape, dtype=float)
        p = 1
    else:
        if len(X.shape) == 1:
            p = 2
        else:
            p = X.shape[1] + 1

    model = sm.GLM(endog=y.astype(float), exog=sm.add_constant(X.astype(float)), family=family)
    result = model.fit()
    resid = result.resid_response

    """Squared standard error of the parameters (i.e. Beta_se)"""
    s2 = (1. / (n-p)) * np.sum(resid**2.)

    np.random.seed(seed)
    if isinstance(Ks, list):
        """If there are multiple kernels then use a min-P omnibus test"""
        Ks = [K.values if isinstance(K, pd.DataFrame) else K for K in Ks]

        obsQ = np.array([computeQ(K, resid, s2) for K in Ks])[None,:]
        randQ = np.nan * np.zeros((nperms, len(Ks)))
        for permi in range(nperms):
            rind = np.random.permutation(n)
            randQ[permi,:] = [computeQ(K[:, rind][rind,:], resid, s2) for K in Ks]
        
        Qall = np.concatenate((obsQ, randQ), axis=0)
        pall = np.zeros(Qall.shape)
        for qi in range(len(Ks)):
            pall[:, qi] = 1 - argrank(Qall[:, qi]) / (nperms + 1.)
        pIndivid = pall[0,:]
        minPall= pall.min(axis=1)
        pGlobal = argrank(minPall)[0] / (nperms + 1.)
        return pGlobal, pIndivid
    else:
        if isinstance(Ks, pd.DataFrame):
            Ks = Ks.values

        obsQ = computeQ(Ks, resid, s2)
        randQ = np.nan * np.zeros(nperms)
        for permi in range(nperms):
            rind = np.random.permutation(n)
            randQ[permi] = computeQ(Ks[:, rind][rind,:], resid, s2)
        pvalue = (np.sum(randQ > obsQ) + 1.) / (nperms + 1.)
        if returnPerms:
            return pvalue, obsQ, randQ
        else:
            return pvalue

def computeKregStat(y, K, X=None, binary=True, alpha=None, nstraps=1000):
    """Compute the statistic that is subjected to permutation testing
    in kernel regression of Y adjusting for covariates in X.

    Parameters
    ----------
    y : pd.Series or np.ndarray shape (n,1)
        Endogenous/outcome variable.
        Can be continuous or binary (use binary=True)
    K : pd.DataFrame or np.ndaray shape (n,n)
        Kernel (square, symetric and positive semi-definite)
    X : pd.DataFrame OR np.ndarray shape (n,i)
        Matrix of covariates for adjustment.