class GP2KronSum(object):
    """
    Gaussian Process with a 2kronSum Covariance and a mean that is a sum of Kronecker products:
        vec(Y) ~ N( vec( \sum_i F_i B_i A_i), Cg \kron R + Cn \kron I )
    Notation:
        N = number of samples
        P = number of traits
        Y = [N, P] phenotype matrix
        F_i = sample fixed effect design for term i
        A_i = trait fixed effect design for term i
        B_i = effect sizes of fixed effect term i
        Cg = column covariance matrix for signal term respectively
        Cn = column covariance matrix for noise term respectively
        R = row covariance matrix for signal term respectively
    """

    def __init__(self, Y, Cg, Cn, F=None, A=None, R=None):
        """
        Args:
            Y:      [N, P] phenotype matrix
            F:      list of sample fixed effect designs.
                    Each term must have first dimension N
            A:      list of trait fixed effect design.
                    Each term must have second dimension P
            Cg:     Limix covariance matrix for Cg (dimension P)
            Cn:     Limix covariance matrix for Cn (dimension P)
            R:      [N, N] numpy semidemidefinite covariance matrix for R.
                    In alternative to R, S_R and U_R can be specified.
            S_R:    N vector of eigenvalues of R
            U_R:    [N, N] eigenvector matrix of R
        """
        assert_type(Y, NP.ndarray, 'Y')
        assert_subtype(Cg, Covariance, 'Cg')
        assert_subtype(Cn, Covariance, 'Cn')

        covar = Cov2KronSum(Cg=Cg, Cn=Cn, R=R, S_R=S_R, U_R=U_R)
        mean = MeanKronSum(Y=Y, F=F, A=A)

        GP.__init__(self, covar=covar, mean=mean)

    # def _observe(self):
    #     self.covar.register(self.row_cov_has_changed, 'row_cov')
    #     self.covar.register(self.col_cov_has_changed, 'col_cov')
    #     self.mean.register(self.pheno_has_changed, 'pheno')
    #     self.mean.register(self.designs_have_changed, 'designs')
    #
    # def col_cov_has_changed(self):
    #     self.clear_cache('col_cov')
    #     self.clear_all()
    #
    # def row_cov_has_changed(self):
    #     self.clear_cache('row_cov')
    #     self.clear_all()
    #
    # def pheno_has_changed(self):
    #     self.clear_cache('pheno')
    #     self.clear_all()
    #
    # def designs_have_changed(self):
    #     self.clear_cache('designs')
    #     self.clear_all()
    #
    #
    # ######################
    # # Transformed phenotype
    # ######################
    # @cached('row_cov', 'pheno')
    # def LrY(self):
    #     return sp.dot(self.covar.Lr(), self.mean.Y)
    #
    # @cached(['row_cov', 'col_cov', 'pheno'])
    # def LrYLc(self):
    #     return sp.dot(self.LrY(), self.covar.Lc().T)
    #
    # @cached(['row_cov', 'col_cov', 'pheno'])
    # def DLrYLc(self):
    #     return self.covar.D() * self.LrYLc()
    #
    # ######################
    # # Transformed fixed effects
    # ######################
    # @cached(['row_cov', 'designs'])
    # def LrF(self):
    #     R = []
    #     for ti in range(self.mean.n_terms):
    #         R.append(sp.dot(self.covar.Lr(), self.mean.F[ti]))
    #     return R
    #
    # @cached(['col_cov', 'designs'])
    # def ALc(self):
    #     R = []
    #     for ti in range(self.mean.n_terms):
    #         R.append(sp.dot(self.mean.A[ti], self.covar.Lc().T))
    #     return R
    #
    # @cached(['row_cov', 'col_cov', 'designs'])
    # def LW(self):
    #     R = sp.zeros((self.mean.Y.size, self.mean.n_covs))
    #     istart = 0
    #     for ti in range(self.mean.n_terms):
    #         iend = istart + self.mean.F[ti].shape[1] * self.mean.A[ti].shape[0]
    #         R[:, istart:iend] = sp.kron(self.ALc()[ti].T, self.LrF()[ti])
    #         istart += iend
    #     return R
    #
    # @cached(['row_cov', 'col_cov', 'designs'])
    # def dLW(self):
    #     return self.covar.d()[:,sp.newaxis] * self.LW()
    #
    # @cached(['row_cov', 'col_cov', 'pheno'])
    # def Sr_DLrYLc_Ctilde(self, i):
    #     return self.covar.Sr_X_Ctilde(self.DLrYLc(), i)
    #
    # @cached(['row_cov', 'col_cov', 'designs', 'pheno'])
    # def Sr_vei_dLWb_Ctilde(self, i):
    #     return self.covar.Sr_X_Ctilde(self.vei_dLWb(), i)
    #
    # @cached(['row_cov', 'col_cov', 'designs', 'pheno'])
    # def vei_dLWb(self):
    #     # could be optimized but probably not worth it
    #     # as it requires a for loop
    #     r = sp.dot(self.dLW(), self.mean.b)
    #     return r.reshape(self.mean.Y.shape, order = 'F')
    #
    # ######################
    # # Areml
    # ######################
    # def Areml_K(self):
    #     return sp.dot(self.LW().T, self.dLW())
    #
    # def Areml_K_grad_i(self,i):
    #     dLWt = self.dLW().reshape((self.mean._N, self.mean._P, self.mean.n_covs), order = 'F')
    #     if i < self.covar.Cg.getNumberParams():
    #         SrdLWt = self.covar.Sr()[:, sp.newaxis, sp.newaxis] * dLWt
    #     else:
    #         SrdLWt = dLWt
    #     SrdLWtC = sp.tensordot(SrdLWt, self.covar.Ctilde(i), axes=(1, 1))
    #     SroCdLW = SrdLWtC.swapaxes(1,2).reshape((self.mean._N * self.mean._P, self.mean.n_covs), order = 'F')
    #     return -sp.dot(self.dLW().T, SroCdLW)
    #
    # ########################
    # # LML terms
    # ########################
    # @cached(['row_cov', 'col_cov', 'designs', 'pheno'])
    # def WKiy(self):
    #     R = sp.zeros((self.mean.n_covs, 1))
    #     istart = 0
    #     for ti in range(self.mean.n_terms):
    #         _dim = self.mean.F[ti].shape[1] * self.mean.A[ti].shape[0]
    #         iend = istart + _dim
    #         FLrDLrYLc = sp.dot(self.LrF()[ti].T, self.DLrYLc())
    #         R[istart:iend, 0] = sp.dot(FLrDLrYLc, self.ALc()[ti].T).reshape(_dim, order = 'F')
    #     return R
    #
    # @cached(['row_cov', 'col_cov', 'designs', 'pheno'])
    # def b(self):
    #     if self.mean.n_covs > 0:
    #         R = self.Areml.solve(self.WKiy())
    #     else:
    #         R = None
    #     return R
    #
    # @cached(['row_cov', 'col_cov', 'pheno'])
    # def yKiy(self):
    #     return (self.LrYLc()*self.DLrYLc()).sum()
    #
    # @cached(['row_cov', 'col_cov', 'designs', 'pheno'])
    # def yKiWb(self):
    #     return (self.LrYLc() * self.vei_dLWb()).sum()
    #
    # #########################
    # # Gradients
    # #########################
    # @cached(['row_cov', 'col_cov', 'pheno'])
    # def yKiy_grad_i(self,i):
    #     return -(self.DLrYLc() * self.Sr_DLrYLc_Ctilde(i)).sum()
    #
    # @cached(['row_cov', 'col_cov', 'designs', 'pheno'])
    # def yKiWb_grad_i(self,i):
    #     rv = -2*(self.DLrYLc()*self.Sr_vei_dLWb_Ctilde(i)).sum()
    #     rv+= (self.vei_dLWb()*self.Sr_vei_dLWb_Ctilde(i)).sum()
    #     return rv
    #
    # ############################
    # # Simulate
    # ############################
    # def simulate_pheno(self, Rh=None):
    #     if Rh is None:
    #         S, U = la.eigh(self.covar.R)
    #         S[S<0] = 0
    #         Rh = U*S**(0.5)
    #     Yc = sp.dot(self.mean.F[0], sp.dot(self.mean.B[0], self.mean.A[0].T))
    #     Z = sp.randn(Rh.shape[1], self.covar.Cr.X.shape[1])
    #     Yr = sp.dot(Rh, sp.dot(Z, self.covar.Cr.X.T))
    #     _S, _U = LA.eigh(self.covar.Cn.K()); _S[_S<0] = 0
    #     Cn_h = _U*_S**0.5
    #     Yn = sp.dot(sp.randn(*self.mean.Y.shape), Cn_h.T)
    #     RV = Yc+Yr+Yn
    #     return RV
