# -*- coding: utf-8 -*-
import numpy as np
import scipy.stats as spst


def make_stat_computer_nopattern(cov, stat_func, **kwargs):
    """
    Create the function that computes the joint statistics
    if no NaN value is in z.

    :param cov: covariance matrix
    :type cov: pandas.core.frame.DataFrame
    :param stat_func: function that computes the joint statistics
    :type stat_func: function
    :return: the function used to compute joint statistics with z as input
    :rtype: function
    """
    # invcov is only computed once
    invcov = np.linalg.pinv(cov, rcond=0.001)#np.linalg.inv(cov)

    def compute(z):
        return stat_func(z, cov, invcov,**kwargs)
    return compute


def make_stat_computer_pattern(cov, stat_func):
    """
    Create the function that computes the joint statistics
    if NaN values are in z. It uses a covariance matrix corresponding
    to the pattern of non-NaN values in z.

    :param cov: covariance matrix
    :type cov: pandas.core.frame.DataFrame
    :param stat_func: function that computes the joint statistics
    :type stat_func: function
    :return: the function used to compute joint statistics with z as input
    :rtype: function
    """
    if not stat_func.can_use_pattern:
        raise ValueError("this computation strategy cannot be used with patterns")
    # invcov_bypattern is a dictionary of invcovs where the key is the
    # corresponding pattern of non-NaN values in z
    invcov_bypattern = {}

    def compute(z):
        z_na_bool = 1 - z.iloc[0,].isnull()
        pattern_code = np.dot(z_na_bool, 10 ** np.arange((len(z_na_bool) - 1), -1, -1))
        z_na_bool = z_na_bool.astype(bool)
        if pattern_code in invcov_bypattern:
            invcov = invcov_bypattern[pattern_code]
        else:
            mini_cov = cov.loc[z_na_bool, z_na_bool]
            invcov = np.linalg.pinv(mini_cov, rcond=0.001)
            invcov_bypattern[pattern_code] = invcov
        z = z.loc[:, z_na_bool]

        return stat_func(z, None, invcov)

    return compute


def make_stat_computer_nan_dumb(cov, stat_func):
    # invert covariance for each line
    if not (stat_func.can_use_pattern):
        raise ValueError("this computation strategy cannot be used with patterns")

    def compute(z):
        z_na_pattern = np.abs((np.isnan(z)).astype(int) - 1)
        z_na_bool = [bool(val) for val in z_na_pattern]
        mini_cov = cov.loc[z_na_bool, z_na_bool]
        invcov = np.linalg.pinv(mini_cov, rcond=0.001)
        z = z.dropna()
        return stat_func(z, None, invcov)

    return compute

def omnibus_stat(z, cov, invcov):
    """
    joint statistics "omnibus" strategy

    note that the omnibus statistics uses invcov but not cov,
    all statistics use the same signature as a simple way to
    allow extensibility

    omnibus can be used in the "NaN processing" function

    :param z: z-scores for the phenotypes
    :type z: pandas.core.series.Series
    :param cov: covariance matrix
    :type cov: DataFrame
    :param invcov: inverted covariance matrix
    :type invcov: numpy.ndarray
    :return: the joint statistics
    :rtype: float
    """
    try:
        p = np.linalg.matrix_rank(invcov)
        #print("DF_K: {0} DF_K': {1}".format(invcov.shape[0], p))
        # stat = np.sum(np.multiply(z, z.dot(invcov)))
        stat = np.einsum("ij,jk,ki->i", z, invcov, z.T)
        return spst.chi2.sf(stat, df=p)
    except ValueError:
        print(z.head())
        print(invcov.shape)
        print("Error in omnibus stat")


omnibus_stat.can_use_pattern = True


def fisher_test(z, cov, invcov):
    """
    fisher combination of univariate p-values corrected with bonferoni

    This test is potentially strongly conservative if the phenotypes are strongly correlated

    :param z: z-scores for the phenotypes
    :type z: pandas.core.series.Series
    :param cov: covariance matrix
    :type cov: DataFrame
    :param invcov: inverted covariance matrix
    :type invcov: numpy.ndarray
    :return: the joint statistics
    :rtype: float
    """
    try:
        p = (~z.isnull()).sum(1)
        print(z)
        print(cov)
        print(invcov)
        p_val = 2*spst.norm.sf(np.abs(z))
        stat = -2 * np.log(np.nansum(p_val, axis=1))
        return spst.chi2.sf(stat, df=p)
    except ValueError:
        print(z.head())
        print(invcov.shape)
        print("Error in Fisher stat")

fisher_test.can_use_pattern = False

def meta_analysis(z, cov, invcov, **kwargs):
    """
    Meta analysis using global sample size to weight z-score

    :param z: z-scores matrix for the phenotypes ((N_snp,N_phenotype) format)
    :type z: pandas.core.DataFrame
    :param cov: covariance matrix
    :type cov: DataFrame
    :param invcov: inverted covariance matrix
    :type invcov: numpy.ndarray
    :return: the joint statistics
    :rtype: numpy.ndarray float64
    :param samp_size: Number of sample (will be used as weights in the meta analysis)
    :type samp_size : pandas.Series
    """

    Effective_sample_size = kwargs.get('samp_size', None)
    if Effective_sample_size is None:
        raise Error('no sample size available to perform meta_analysis')
    else:
        loading = Effective_sample_size.loc[z.columns]**0.5

        M_loadings = np.full(z.shape, loading**2)
        M_loadings[np.isnan(z)] = 0

        z = np.nan_to_num(z) # fill na with zero

        numi = loading.dot(z.transpose())
        deno = np.sqrt(np.sum(M_loadings, axis=1))
        # print(loading)
        # print(z[:10,:])
        # print(numi[:10])
        # print(deno[:10])
        # print(stat[:10])
        # fill na with 0 = don't take the missing GWAS into account in the test
        stat = numi / deno

        return spst.chi2.sf(stat, df=1)

meta_analysis.can_use_pattern = False

def sumz_stat(z, cov, invcov, **kwargs):
    """
    joint statistics "sumZ" strategy

    :param z: z-scores matrix for the phenotypes ((N_snp,N_phenotype) format)
    :type z: pandas.core.DataFrame
    :param cov: covariance matrix
    :type cov: DataFrame
    :param invcov: inverted covariance matrix
    :type invcov: numpy.ndarray
    :return: the joint statistics
    :rtype: numpy.ndarray float64
    """
    loading = kwargs.get('loadings', None)

    if loading is None:
        p = z.shape[1]
        loading = np.ones(p)
    else:
        loading = loading.loc[z.columns]

    print(loading)
    print(cov)
    M_loadings = np.full(z.shape, loading)
    M_loadings[np.isnan(z)] = 0

    z = np.nan_to_num(z)
    numi = np.square(loading.dot(z.transpose()))
    deno = np.einsum('ij,jk,ki->i', M_loadings, cov, M_loadings.T)

    # fill na with 0 = don't take the missing GWAS into account in the test
    stat = numi / deno
    return spst.chi2.sf(stat, df=1)


sumz_stat.can_use_pattern = False
