import numpy as NP
import numpy as np
import scipy.linalg as LA
import scipy.optimize as opt
import scipy.stats as ST
import scipy.special as SS
from fastlmm.util.mingrid import *
from fastlmm.util.util import *
import time

from bin2kernel import Bin2Kernel
from bin2kernel import makeBin2KernelAsEstimator
from bin2kernel import Bin2KernelLaplaceLinearN
from bin2kernel import getFastestBin2Kernel
from bin2kernel import Bin2KernelEPLinearN


def rotate(A, rotationMatrix, transposeRot=True):
    """
    Apply an eigenvalue decomposition rotation matrix (Assummes that the rotation-matrix has orthogonal cols)

    takes care of low rank structure in A

    result = eig[1].T * A
    """
    [Nr, k] = rotationMatrix.shape
    [N, M] = A.shape
    if Nr != N:
        raise Exception(
            "rotation and A are not aligned A.shape =  [%i,%i], rotation.shape = [%i,%i]"
            % (N, M, Nr, k)
        )
    if k > Nr:
        raise Exception(
            "rotation matrix has more columns than rows =  [%i,%i], rotation.shape = [%i,%i]"
            % (N, M, Nr, k)
        )

    res = rotationMatrix.T.dot(A)

    if k < Nr:
        resLowRank = A - rotationMatrix.dot(res)
        res = [res, resLowRank]
    return res


def rotSymm(A, eig, delta=1.0, gamma=1.0, exponent=-0.5, forceSymm=True):
    """
    do the rotation with K_0 to get a matrix square root
    """
    [N, M] = A.shape
    [Nr, k] = eig[1].shape
    if Nr != N:
        raise Exception(
            "rotation and A are not aligned A.shape =  [%i,%i], rotation.shape = [%i,%i]"
            % (N, M, Nr, k)
        )
    if symmetric and N != M:
        raise Exception(
            "A is not symmetric. A.shape =  [%i,%i], rotation.shape = [%i,%i]"
            % (N, M, Nr, k)
        )
    if k > Nr:
        raise Exception(
            "rotation matrix has more columns than rows =  [%i,%i], rotation.shape = [%i,%i]"
            % (N, M, Nr, k)
        )

    res = eig[1].T.dot(A)
    deltaPow = myPow(delta, exponent, minVal)
    diag = eig[0] * gamma + delta
    diag = myPow(diag, exponent, minVal=0.0)
    if Nr > k:
        diag -= deltaPow
        res = dotDiag(res, diag)
        res = A - eig[1].dot(res)
    else:
        res = dotDiag(res, diag)
        if forceSymm:
            res = eig[1].dot(res)
    return res


def dotDiag(A, diag, minVal=None, symmetric=False, exponent=1.0):
    """
    Multiply by a diagonal matrix build from

    result = diag.dot(A)
    """
    [N, M] = A.shape
    [Nd] = diag.shape
    if symmetric and N != M:
        raise Exception(
            "A is not symmetric. A.shape =  [%i,%i], rotation.shape = [%i,%i]"
            % (N, M, Nr, k)
        )
    if Nd != N:
        raise Exception(
            "Matrices misaligned A.shape =  [%i,%i], diag.shape = [%i,%i]"
            % (N, M, Nr, k)
        )

    diagcorr = myPow(diag, exponent, minVal)

    result = A * NP.lib.stride_tricks.as_strided(
        diagcorr, (diagcorr.size, A.shape[1]), (diagcorr.itemsize, 0)
    )
    if symmetric:
        result = dotDiag(
            A=A.T,
            diag=diagcorr,
            checkPos=False,
            minVal=minVal,
            symmetric=False,
            exponent=1.0,
        ).T
    return result


def myPow(A, exponent, minVal=None):
    """
    compute a power, with efficient computation of common powers
    """
    if exponent == 1.0:
        Aexp = A
    elif exponent == -1.0:
        Aexp = 1.0 / A
    elif exponent == 2.0:
        Aexp = A * A
    elif exponent == 0.5:
        Aexp = NP.sqrt(A)
    elif exponent == -0.5:
        Aexp = 1.0 / NP.sqrt(A)
    elif exponent == -2.0:
        Aexp = 1.0 / (A * A)
    else:
        Aexp = NP.power(A, exponent)
    if minVal is not None:
        i_Null = A <= minVal
        Aexp[i_Null] = 0.0
    return Aexp


class lmm2k(object):
    """
    linear mixed model with up to two kernels
    N(y | X*beta ; sigma2(gamma0*K0 + gamma1*K1 + I ) ),
    where
    K0 = G0*G0^T
    K1 = G1*G1^T
    """

    __slots__ = [
        "G0",
        "G1",
        "Y",
        "X",
        "K0",
        "K1",
        "K",
        "U",
        "S",
        "UX",
        "Uy",
        "UUX",
        "UW",
        "UUW",
        "UUy",
        "pos0",
        "pos1",
        "gamma0",
        "gamma1",
        "delta",
        "exclude_idx",
        "forcefullrank",
        "numcalls",
        "Xstar",
        "Kstar",
        "Kstar_star",
        "UKstar",
        "UUKstar",
        "Gstar",
    ]

    def __init__(self, forcefullrank=False):
        """
        Input:
        forcefullrank   : if True, then the code always computes K and runs cubically
                            (False)
        """
        self.X = None
        self.Y = None
        self.G0 = None
        self.G1 = None
        self.K0 = None
        self.K1 = None
        self.gamma0 = 1.0
        self.gamma1 = 1.0
        self.delta = 1.0
        self.eig0 = None
        self.eig1 = None
        self.Yrot = None
        self.Xrot = None
        self.K1rot = None
        self.G1rot = None

    def setX(self, X):
        """
        set the fixed effects X (covariates).
        The Kernel has to be set in advance by first calling setG() or setK().
        --------------------------------------------------------------------------
        Input:
        X       : [N*D] 2-dimensional array of covariates
        --------------------------------------------------------------------------
        """
        self.X = X
        self.Xrot = None

    def setY(self, Y):
        """
        set the phenotype y.
        The Kernel has to be set in advance by first calling setG() or setK().
        --------------------------------------------------------------------------
        Input:
        Y       : [NxP] P-dimensional array of phenotype values
        --------------------------------------------------------------------------
        """
        self.Y = Y
        self.Yrot = None

    def setG0(self, G0):
        """
        set the Kernel K0 from G0.
        This has to be done before setting the data setX() and setY().
        ----------------------------------------------------------------------------
        Input:
        G0              : [N*k0] array of random effects
        -----------------------------------------------------------------------------
        """
        k = G0.shape[1]
        N = G0.shape[0]
        self.eig0 = None
        self.eig1 = None
        self.K1rot = None
        self.Yrot = None
        self.Xrot = None

        if (not self.forcefullrank) and (k < N):
            self.G0 = G0
        else:
            K0 = G0.dot(G0.T)
            self.setK0(K0=K0)

    def setK0(self, K0):
        """
        set the background Kernel K0.
        --------------------------------------------------------------------------
        Input:
        K0 : [N*N] array, random effects covariance (positive semi-definite)
        --------------------------------------------------------------------------
        """
        self.K0 = K0
        self.G1 = None
        self.eig0 = None
        self.eig1 = None
        self.K1rot = None
        self.Yrot = None
        self.Xrot = None
        self.G1rot = None
        self.K1rot = None

    def setG1(self, G1):
        """
        set the Kernel K1 from G1.
        ----------------------------------------------------------------------------
        Input:
        G0              : [N*k0] array of random effects
        -----------------------------------------------------------------------------
        """
        k = self.G1.shape[1]
        N = self.G0.shape[0]
        self.eig1 = None
        self.G1rot = None
        if (not self.forcefullrank) and (k < N):
            # it is faster using the eigen decomposition of G.T*G but this is more accurate
            self.G1 = G1
        else:
            K1 = self.G1.dot(self.G1.T)
            self.setK1(K1=K1)
        pass

    def setK1(self, K1):
        """
        set the foreground Kernel K1.
        This has to be done before setting the data setX() and setY().
        --------------------------------------------------------------------------
        Input:
        K1 : [N*N] array, random effects covariance (positive semi-definite)
        --------------------------------------------------------------------------
        """
        self.eig1 = None
        self.G1 = None
        self.K1 = K1
        self.K1rot = None
        self.G1rot = None

    def setVariances(self, gamma0=None, gamma1=None, delta=None):
        if gamma0 is not None:
            # background model changed, foreground rotation (U1,S1) changes:
            self.gamma0 = gamma0
            self.eig1 = None
            self.Xrot = None
            self.Yrot = None
            self.K1rot = None
            self.G1rot = None

        if delta is not None:
            # background model changed, foreground rotation (U1,S1) changes:
            self.delta = delta
            self.eig1 = None
            self.Xrot = None
            self.Yrot = None
            self.K1rot = None
            self.G1rot = None

        if gamma1 is not None:
            # foreground model changed
            self.gamma1 = gamma1

    def getEig1(self):
        """ """
        if self.eig1 is None:
            # compute eig1
            if self.K1 is not None:
                if self.K1rot is None:
                    self.K1rot = rotSymm(
                        self.K1,
                        eig=self.eig0,
                        exponent=-0.5,
                        gamma=self.gamma0,
                        delta=self.delta,
                        forceSymm=False,
                    )
                    self.K1rot = rotSymm(
                        self.K1rot.T,
                        eig=self.eig0,
                        exponent=-0.5,
                        gamma=self.gamma0,
                        delta=self.delta,
                        forceSymm=False,
                    )
                self.eig1 = LA.eigh(self.K1rot)
            elif self.G1 is not None:
                [N, k] = self.G1.shape
                if self.G1rot is None:
                    self.G1rot = rotSymm(
                        self.G1,
                        eig=self.eig0,
                        exponent=-0.5,
                        gamma=self.gamma0,
                        delta=self.delta,
                        forceSymm=False,
                    )

                try:
                    [U, S, V] = LA.svd(self.G1rot, full_matrices=False)
                    self.eig1 = [S * S, U]
                except LA.LinAlgError:  # revert to Eigenvalue decomposition
                    print(
                        "Got SVD exception, trying eigenvalue decomposition of square of G. Note that this is a little bit less accurate"
                    )
                    [S_, V_] = LA.eigh(self.G1rot.T.dot(self.G1rot))
                    S_nonz = S_ > 0.0
                    S1 = S_[S_nonz]
                    U1 = self.G1rot.dot(V_[:, S_nonz] / np.sqrt(S1))
                    self.eig1 = [S1, U1]
        return self.eig1

    def getEig0(self):
        """ """
        if self.eig0 is None:
            # compute eig0
            if self.K0 is not None:
                self.eig1 = LA.eigh(self.K0)
            elif self.G1 is not None:
                [N, k] = self.G0.shape
                try:
                    [U, S, V] = LA.svd(self.G0, full_matrices=False)
                    self.eig0 = [S * S, U]
                except LA.LinAlgError:  # revert to Eigenvalue decomposition
                    print(
                        "Got SVD exception, trying eigenvalue decomposition of square of G. Note that this is a little bit less accurate"
                    )
                    [S_, V_] = LA.eigh(self.G0.T.dot(self.G0))
                    S_nonz = S_ > 0.0
                    S0 = S_[S_nonz]
                    U0 = self.G0.dot(V_[:, S_nonz] / np.sqrt(S0))
                    self.eig0 = [S0, U0]
        return self.eig1

    def set_exclude_idx(self, idx):
        """
        Set the indices of SNPs to be removed
        --------------------------------------------------------------------------
        Input:
        idx  : [k_up: number of SNPs to be removed] holds the indices of SNPs to be removed
        --------------------------------------------------------------------------
        """
        self.exclude_idx = idx

    def findGamma1givenGamma0(
        self, gamma0=0.0, nGridGamma1=10, minGamma1=0.0, maxGamma1=10000.0, **kwargs
    ):
        """
        Find the optimal h2 for a given K. Note that this is the single kernel case. So there is no a2.
        (default maxH2 value is set to a value smaller than 1 to avoid loss of positive definiteness of the final model covariance)
        --------------------------------------------------------------------------
        Input:
        nGridH2 : number of h2-grid points to evaluate the negative log-likelihood at
        minH2   : minimum value for h2 optmization
        maxH2   : maximum value for h2 optmization
        --------------------------------------------------------------------------
        Output:
        dictionary containing the model parameters at the optimal h2
        --------------------------------------------------------------------------
        """
        self.setVariances(gamma0=gamma0)
        resmin = [None]

        def f(x, resmin=resmin, **kwargs):
            # self.setVariances(gamma1=x)
            res = self.nLLeval(gamma1=x, **kwargs)
            if (resmin[0] is None) or (res["nLL"] < resmin[0]["nLL"]):
                resmin[0] = res
                self.setGamma1[x]
            return res["nLL"]

        min = minimize1D(f=f, nGrid=nGridGamma1, minval=minGamma1, maxval=maxGamma1)
        return resmin[0]

    def findGammas(
        self,
        nGridGamma0=10,
        minGamma0=0.0,
        maxGamma0=10000.0,
        nGridGamma1=10,
        minGamma1=0.0,
        maxGamma1=10000.0,
        verbose=False,
        **kwargs
    ):
        """
        Find the optimal gamma0 and gamma1, such that K=gamma0*K0+gamma1*K1. Performs a double loop optimization (could be expensive for large grid-sizes)
        --------------------------------------------------------------------------
        Input:
        nGridGamma0 : number of a2-grid points to evaluate the negative log-likelihood at
        minGamma0   : minimum value for a2 optmization
        maxGamma0   : maximum value for a2 optmization
        nGridGamma1 : number of h2-grid points to evaluate the negative log-likelihood at
        minGamma1   : minimum value for h2 optmization
        maxGamma1   : maximum value for h2 optmization
        --------------------------------------------------------------------------
        Output:
        dictionary containing the model parameters at the optimal gamma0 and gamma1
        --------------------------------------------------------------------------
        """
        self.numcalls = 0
        resmin = [None]

        def f(
            x,
            resmin=resmin,
            nGridGamma1=nGridGamma1,
            minGamma1=minGamma1,
            maxGamma1=maxGamma1,
            **kwargs
        ):
            self.numcalls += 1
            t0 = time.time()
            res = self.findGamma1givenGamma0(
                gamma0=x,
                nGridGamma1=nGridGamma1,
                minGamma1=minGamma1,
                maxGamma1=maxGamma1,
                **kwargs
            )
            if (resmin[0] is None) or (res["nLL"] < resmin[0]["nLL"]):
                resmin[0] = res
            t1 = time.time()
            # print "one objective function call took %.2f seconds elapsed" % (t1-t0)
            # import pdb; pdb.set_trace()
            return res["nLL"]

        if verbose:
            print("findGammas")
        min = minimize1D(
            f=f, nGrid=nGridGamma1, minval=minGamma1, maxval=maxGamma1, verbose=False
        )
        # print "numcalls to innerLoopTwoKernel= " + str(self.numcalls)
        return resmin[0]

    def nLLeval(self, REML=True, gamma1=None, dof=None, scale=1.0):
        """
        evaluate -ln( N( U^T*y | U^T*X*beta , h2*S + (1-h2)*I ) ),
        where ((1-a2)*K0 + a2*K1) = USU^T
        --------------------------------------------------------------------------
        Input:
        h2      : mixture weight between K and Identity (environmental noise)
        REML    : boolean
                  if True   : compute REML
                  if False  : compute ML
        dof     : Degrees of freedom of the Multivariate student-t
                        (default None uses multivariate Normal likelihood)
        logdelta: log(delta) allows to optionally parameterize in delta space
        delta   : delta     allows tomoptionally parameterize in delta space
        scale   : Scale parameter the multiplies the Covariance matrix (default 1.0)
        --------------------------------------------------------------------------
        Output dictionary:
        'nLL'       : negative log-likelihood
        'sigma2'    : the model variance sigma^2
        'beta'      : [D*1] array of fixed effects weights beta
        'h2'        : mixture weight between Covariance and noise
        'REML'      : True: REML was computed, False: ML was computed
        'a2'        : mixture weight between K0 and K1
        'dof'       : Degrees of freedom of the Multivariate student-t
                        (default None uses multivariate Normal likelihood)
        'scale'     : Scale parameter that multiplies the Covariance matrix (default 1.0)
        --------------------------------------------------------------------------
        """
        if gamma1 is None:
            gamma1 = self.gamma1
        if gamma1 < 0.0:
            return {
                "nLL": 3e20,
                "gamma1": gamma1,
                "delta": self.delta,
                "gamma0": self.gamma0,
                "dof": dof,
                "REML": REML,
                "scale": scale,
            }
        k = self.S1.shape[0]
        N = self.Y.shape[0]
        D = self.X.shape[1]

        Sd = (gamma1 * self.S1 + 1.0) * scale
        UXS = dotDiag(self.UX, Sd, minVal=0.0, symmetric=False, exponent=-1.0)
        UYS = dotDiag(self.UY, Sd, minVal=0.0, symmetric=False, exponent=-1.0)

        XKX = UXS.T.dot(self.UX)
        XKy = UXS.T.dot(self.Uy)
        yKy = UyS.T.dot(self.Uy)

        logdetK = np.log(Sd).sum()

        if k < N:  # low rank part
            # determine normalization factor
            denom = 1.0 * scale

            XKX += self.UUX.T.dot(self.UUX) / (denom)
            XKy += self.UUX.T.dot(self.UUy) / (denom)
            yKy += self.UUy.T.dot(self.UUy) / (denom)
            logdetK += (N - k) * np.log(denom)

        # proximal contamination (see Supplement Note 2: An Efficient Algorithm for Avoiding Proximal Contamination)
        # available at: http://www.nature.com/nmeth/journal/v9/n6/extref/nmeth.2037-S1.pdf
        # exclude SNPs from the RRM in the likelihood evaluation

        if len(self.exclude_idx) > 0:
            raise Exception("not implemented")
            num_exclude = len(self.exclude_idx)

            # consider only excluded SNPs
            G_exclude = self.G[:, self.exclude_idx]

            self.UW = self.U.T.dot(G_exclude)  # needed for proximal contamination
            UWS = self.UW / NP.lib.stride_tricks.as_strided(
                Sd, (Sd.size, num_exclude), (Sd.itemsize, 0)
            )
            assert UWS.shape == (k, num_exclude)

            WW = NP.eye(num_exclude) - UWS.T.dot(self.UW)
            WX = UWS.T.dot(self.UX)
            Wy = UWS.T.dot(self.Uy)
            assert WW.shape == (num_exclude, num_exclude)
            assert WX.shape == (num_exclude, D)
            assert Wy.shape == (num_exclude,)

            if k < N:  # low rank part
                self.UUW = G_exclude - self.U.dot(self.UW)

                WW += self.UUW.T.dot(self.UUW) / denom
                WX += self.UUW.T.dot(self.UUX) / denom
                Wy += self.UUW.T.dot(self.UUy) / denom

            # TODO: do cholesky, if fails do eigh
            # compute inverse efficiently
            [S_WW, U_WW] = LA.eigh(WW)

            UWX = U_WW.T.dot(WX)
            UWy = U_WW.T.dot(Wy)
            assert UWX.shape == (num_exclude, D)
            assert UWy.shape == (num_exclude,)

            # compute S_WW^{-1} * UWX
            WX = UWX / NP.lib.stride_tricks.as_strided(
                S_WW, (S_WW.size, UWX.shape[1]), (S_WW.itemsize, 0)
            )
            # compute S_WW^{-1} * UWy
            Wy = UWy / S_WW
            # determinant update
            logdetK += np.log(S_WW).sum()
            assert WX.shape == (num_exclude, D)
            assert Wy.shape == (num_exclude,)

            # perform updates (instantiations for a and b in Equation (1.5) of Supplement)
            yKy += UWy.T.dot(Wy)
            XKy += UWX.T.dot(Wy)
            XKX += UWX.T.dot(WX)

        #######

        [SxKx, UxKx] = LA.eigh(XKX)
        i_pos = SxKx > 1e-10
        beta = np.dot(UxKx[:, i_pos], (np.dot(UxKx[:, i_pos].T, XKy) / SxKx[i_pos]))

        r2 = yKy - XKy.dot(beta)

        if dof is None:  # Use the Multivariate Gaussian
            if REML:
                XX = self.X.T.dot(self.X)
                [Sxx, Uxx] = LA.eigh(XX)
                logdetXX = np.log(Sxx).sum()
                logdetXKX = np.log(SxKx).sum()
                sigma2 = r2 / (N - D)
                nLL = 0.5 * (
                    logdetK
                    + logdetXKX
                    - logdetXX
                    + (N - D) * (np.log(2.0 * np.pi * sigma2) + 1)
                )
            else:
                sigma2 = r2 / (N)
                nLL = 0.5 * (logdetK + N * (np.log(2.0 * np.pi * sigma2) + 1))
            result = {
                "nLL": nLL,
                "sigma2": sigma2,
                "beta": beta,
                "gamma1": gamma1,
                "REML": REML,
                "gamma0": self.gamma0,
                "delta": self.delta,
                "scale": scale,
            }
        else:  # Use multivariate student-t
            if REML:
                XX = self.X.T.dot(self.X)
                [Sxx, Uxx] = LA.eigh(XX)
                logdetXX = np.log(Sxx).sum()
                logdetXKX = np.log(SxKx).sum()

                nLL = 0.5 * (
                    logdetK
                    + logdetXKX
                    - logdetXX
                    + (dof + (N - D)) * np.log(1.0 + r2 / dof)
                )
                nLL += (
                    0.5 * (N - D) * np.log(dof * np.pi)
                    + SS.gammaln(0.5 * dof)
                    - SS.gammaln(0.5 * (dof + (N - D)))
                )
            else:
                nLL = 0.5 * (logdetK + (dof + N) * np.log(1.0 + r2 / dof))
                nLL += (
                    0.5 * N * np.log(dof * np.pi)
                    + SS.gammaln(0.5 * dof)
                    - SS.gammaln(0.5 * (dof + N))
                )
            result = {
                "nLL": nLL,
                "dof": dof,
                "sigma2": sigma2,
                "beta": beta,
                "gamma1": gamma1,
                "REML": REML,
                "gamma0": self.gamma0,
                "delta": self.delta,
                "scale": scale,
            }

        return result
