# Path: brainiak/matnormal/utils.py
# def rmn(rowcov, colcov):
#     """
#     Generate random draws from a zero-mean matrix-normal distribution.
# 
#     Parameters
#     -----------
#     rowcov : np.ndarray
#         Row covariance (assumed to be positive definite)
#     colcov : np.ndarray
#         Column covariance (assumed to be positive definite)
#     """
# 
#     Z = norm.rvs(size=(rowcov.shape[0], colcov.shape[0]))
#     return cholesky(rowcov).dot(Z).dot(cholesky(colcov))
# 
# Path: brainiak/matnormal/matnormal_likelihoods.py
# def matnorm_logp(x, row_cov, col_cov):
#     """Log likelihood for centered matrix-variate normal density.
#     Assumes that row_cov and col_cov follow the API defined in CovBase.
# 
#     Parameters
#     ----------------
#     x: tf.Tensor
#         Observation tensor
#     row_cov: CovBase
#         Row covariance implementing the CovBase API
#     col_cov: CovBase
#         Column Covariance implementing the CovBase API
# 
#     """
# 
#     rowsize = tf.cast(tf.shape(input=x)[0], "float64")
#     colsize = tf.cast(tf.shape(input=x)[1], "float64")
# 
#     # precompute sigma_col^{-1} * x'
#     solve_col = col_cov.solve(tf.transpose(a=x))
#     logdet_col = col_cov.logdet
# 
#     # precompute sigma_row^{-1} * x
#     solve_row = row_cov.solve(x)
#     logdet_row = row_cov.logdet
# 
#     return _mnorm_logp_internal(
#         colsize, rowsize, logdet_row, logdet_col, solve_row, solve_col
#     )
# 
# Path: brainiak/matnormal/covs.py
# class CovIdentity(CovBase):
#     """Identity noise covariance.
#     """
# 
#     def __init__(self, size):
#         super(CovIdentity, self).__init__(size)
# 
#     @property
#     def logdet(self):
#         return tf.constant(0.0, "float64")
# 
#     def get_optimize_vars(self):
#         """Returns a list of tf variables that need to get optimized to
#             fit this covariance
#         """
#         return []
# 
#     def solve(self, X):
#         """Given this covariance :math:`\\Sigma` and some input :math:`X`,
#             compute :math:`\\Sigma^{-1}x`
#         """
#         return X
# 
#     @property
#     def _prec(self):
#         """Expose the precision explicitly (mostly for testing /
#         visualization, materializing large covariances may be intractable)
#         """
#         return tf.eye(self.size, dtype=tf.float64)
# 
#     @property
#     def _cov(self):
#         """Expose the covariance explicitly (mostly for testing /
#         visualization, materializing large covariances may be intractable)
#         """
#         return tf.eye(self.size, dtype=tf.float64)
# 
# class CovUnconstrainedCholesky(CovBase):
#     """Unconstrained noise covariance parameterized in terms of its cholesky
#     """
# 
#     def __init__(self, size=None, Sigma=None):
# 
#         if size is None and Sigma is None:
#             raise RuntimeError("Must pass either Sigma or size but not both")
# 
#         if size is not None and Sigma is not None:
#             raise RuntimeError("Must pass either Sigma or size but not both")
# 
#         if Sigma is not None:
#             size = Sigma.shape[0]
# 
#         super(CovUnconstrainedCholesky, self).__init__(size)
# 
#         # number of parameters in the triangular mat
#         npar = (size * (size + 1)) // 2
# 
#         if Sigma is None:
#             self.L_flat = tf.Variable(
#                 tf.random.normal([npar], dtype=tf.float64), name="L_flat"
#             )
# 
#         else:
#             L = np.linalg.cholesky(Sigma)
#             self.L_flat = tf.Variable(
#                 flatten_cholesky_unique(L), name="L_flat")
# 
#         self.optimize_vars = [self.L_flat]
# 
#     @property
#     def L(self):
#         """
#         Cholesky factor of this covariance
#         """
#         return unflatten_cholesky_unique(self.L_flat)
# 
#     @property
#     def logdet(self):
#         return 2 * tf.reduce_sum(input_tensor=tf.math.log(
#                                  tf.linalg.diag_part(self.L)))
# 
#     def get_optimize_vars(self):
#         """ Returns a list of tf variables that need to get optimized to fit
#              this covariance
#         """
#         return [self.L_flat]
# 
#     def solve(self, X):
#         """Given this covariance :math:`\\Sigma` and some input :math:`X`,
#         compute :math:`\\Sigma^{-1}x` (using cholesky solve)
# 
#         Parameters
#         ----------
#         X: tf.Tensor
#             Tensor to multiply by inverse of this covariance
# 
#         """
#         return tf.linalg.cholesky_solve(self.L, X)


import numpy as np
import tensorflow as tf
from numpy.testing import assert_allclose
from scipy.stats import multivariate_normal
from brainiak.matnormal.utils import rmn
from brainiak.matnormal.matnormal_likelihoods import matnorm_logp
from brainiak.matnormal.covs import CovIdentity, CovUnconstrainedCholesky


# X is m x n, so A sould be m x p

m = 5
n = 4
p = 3

rtol = 1e-7


def test_against_scipy_mvn_row(seeded_rng):

    rowcov = CovUnconstrainedCholesky(size=m)
    colcov = CovIdentity(size=n)
    X = rmn(np.eye(m), np.eye(n))
    X_tf = tf.constant(X, "float64")

    rowcov_np = rowcov._cov

    scipy_answer = np.sum(multivariate_normal.logpdf(
        X.T, np.zeros([m]), rowcov_np))
    tf_answer = matnorm_logp(X_tf, rowcov, colcov)
    assert_allclose(scipy_answer, tf_answer, rtol=rtol)


def test_against_scipy_mvn_col(seeded_rng):


# Path: brainiak/matnormal/utils.py
# def rmn(rowcov, colcov):
#     """
#     Generate random draws from a zero-mean matrix-normal distribution.
# 
#     Parameters
#     -----------
#     rowcov : np.ndarray
#         Row covariance (assumed to be positive definite)
#     colcov : np.ndarray
#         Column covariance (assumed to be positive definite)
#     """
# 
#     Z = norm.rvs(size=(rowcov.shape[0], colcov.shape[0]))
#     return cholesky(rowcov).dot(Z).dot(cholesky(colcov))
# 
# Path: brainiak/matnormal/matnormal_likelihoods.py
# def matnorm_logp(x, row_cov, col_cov):
#     """Log likelihood for centered matrix-variate normal density.
#     Assumes that row_cov and col_cov follow the API defined in CovBase.
# 
#     Parameters
#     ----------------
#     x: tf.Tensor
#         Observation tensor
#     row_cov: CovBase
#         Row covariance implementing the CovBase API
#     col_cov: CovBase
#         Column Covariance implementing the CovBase API
# 
#     """
# 
#     rowsize = tf.cast(tf.shape(input=x)[0], "float64")
#     colsize = tf.cast(tf.shape(input=x)[1], "float64")
# 
#     # precompute sigma_col^{-1} * x'
#     solve_col = col_cov.solve(tf.transpose(a=x))
#     logdet_col = col_cov.logdet
# 
#     # precompute sigma_row^{-1} * x
#     solve_row = row_cov.solve(x)
#     logdet_row = row_cov.logdet
# 
#     return _mnorm_logp_internal(
#         colsize, rowsize, logdet_row, logdet_col, solve_row, solve_col
#     )
# 
# Path: brainiak/matnormal/covs.py
# class CovIdentity(CovBase):
#     """Identity noise covariance.
#     """
# 
#     def __init__(self, size):
#         super(CovIdentity, self).__init__(size)
# 
#     @property
#     def logdet(self):
#         return tf.constant(0.0, "float64")
# 
#     def get_optimize_vars(self):
#         """Returns a list of tf variables that need to get optimized to
#             fit this covariance
#         """
#         return []
# 
#     def solve(self, X):
#         """Given this covariance :math:`\\Sigma` and some input :math:`X`,
#             compute :math:`\\Sigma^{-1}x`
#         """
#         return X
# 
#     @property
#     def _prec(self):
#         """Expose the precision explicitly (mostly for testing /
#         visualization, materializing large covariances may be intractable)
#         """
#         return tf.eye(self.size, dtype=tf.float64)
# 
#     @property
#     def _cov(self):
#         """Expose the covariance explicitly (mostly for testing /
#         visualization, materializing large covariances may be intractable)
#         """
#         return tf.eye(self.size, dtype=tf.float64)
# 
# class CovUnconstrainedCholesky(CovBase):
#     """Unconstrained noise covariance parameterized in terms of its cholesky
#     """
# 
#     def __init__(self, size=None, Sigma=None):
# 
#         if size is None and Sigma is None:
#             raise RuntimeError("Must pass either Sigma or size but not both")
# 
#         if size is not None and Sigma is not None:
#             raise RuntimeError("Must pass either Sigma or size but not both")
# 
#         if Sigma is not None:
#             size = Sigma.shape[0]
# 
#         super(CovUnconstrainedCholesky, self).__init__(size)
# 
#         # number of parameters in the triangular mat
#         npar = (size * (size + 1)) // 2
# 
#         if Sigma is None:
#             self.L_flat = tf.Variable(
#                 tf.random.normal([npar], dtype=tf.float64), name="L_flat"
#             )
# 
#         else:
#             L = np.linalg.cholesky(Sigma)
#             self.L_flat = tf.Variable(
#                 flatten_cholesky_unique(L), name="L_flat")
# 
#         self.optimize_vars = [self.L_flat]
# 
#     @property
#     def L(self):
#         """
#         Cholesky factor of this covariance
#         """
#         return unflatten_cholesky_unique(self.L_flat)
# 
#     @property
#     def logdet(self):
#         return 2 * tf.reduce_sum(input_tensor=tf.math.log(
#                                  tf.linalg.diag_part(self.L)))
# 
#     def get_optimize_vars(self):
#         """ Returns a list of tf variables that need to get optimized to fit
#              this covariance
#         """
#         return [self.L_flat]
# 
#     def solve(self, X):
#         """Given this covariance :math:`\\Sigma` and some input :math:`X`,
#         compute :math:`\\Sigma^{-1}x` (using cholesky solve)
# 
#         Parameters
#         ----------
#         X: tf.Tensor
#             Tensor to multiply by inverse of this covariance
# 
#         """
#         return tf.linalg.cholesky_solve(self.L, X)


# Path: tests/matnormal/test_matnormal_logp.py
import numpy as np
import tensorflow as tf
from numpy.testing import assert_allclose
from scipy.stats import multivariate_normal
from brainiak.matnormal.utils import rmn
from brainiak.matnormal.matnormal_likelihoods import matnorm_logp
from brainiak.matnormal.covs import CovIdentity, CovUnconstrainedCholesky



# X is m x n, so A sould be m x p

m = 5
n = 4
p = 3

rtol = 1e-7


def test_against_scipy_mvn_row(seeded_rng):

    rowcov = CovUnconstrainedCholesky(size=m)
    colcov = CovIdentity(size=n)
    X = rmn(np.eye(m), np.eye(n))
    X_tf = tf.constant(X, "float64")

    rowcov_np = rowcov._cov

    scipy_answer = np.sum(multivariate_normal.logpdf(
        X.T, np.zeros([m]), rowcov_np))
    tf_answer = matnorm_logp(X_tf, rowcov, colcov)
    assert_allclose(scipy_answer, tf_answer, rtol=rtol)


def test_against_scipy_mvn_col(seeded_rng):


    rowcov = CovIdentity(size=m)