import numpy as np 
from scipy.stats import multivariate_normal
import pdb


class MultivariateGaussianData:
    """ Class for generating data from a multivariate Gaussian distibution 
    (Z, X) ~ N(mu, Sigma), where Z is considered the latent variable and X the observation, 
    such that Z ~ N(0, I) and X|Z=z ~ N(Wz + b, Sigma_x)
    ------------------
    Parameters:
              dim_x : dimension of X
              dim_z : dimension of Z
    ------------------
    Methods:
            create_dataset : generate n_samples observations X ~ N(mu_x, Sigma_x)
            compute_posterior : return p(Z|X=x), mean and covariance of Z|X=x
            compute_log_likelihood : return mean(log(p(x))) across the given dataset
    """

    def __init__(self, dim_x, dim_z):
        np.random.seed(0)
        self.n_x = dim_x
        self.n_z = dim_z
        
        # z ~ N(0, I):
        self.sigma_z = np.eye(dim_z)
        self.mu_z = np.zeros(dim_z)

        # X|Z=z ~ N(Wz+b, Sigma_x_z)
        self.W = np.random.normal(0, 1, size=(dim_x, dim_z))
        self.b = np.random.normal(0, 1, dim_x)
        aux_sigma = np.random.normal(0, 1, size=(dim_x, dim_x))
        self.sigma_x_z = np.dot(aux_sigma, aux_sigma.T) + 0.5 * np.eye(dim_x)
        
        # Cholesky factor of the covariance of the observed variable X (needed for sampling)
        self.chol_sigma_x_z = np.linalg.cholesky(self.sigma_x_z)

        # X + N(b, W*W^T + Sigma_x, z)
        self.sigma_x = np.dot(self.W, self.W.T) + self.sigma_x_z
   

    def create_dataset(self, n_samples):
        """ Sample n_samples from X ~ N(mu_x, Sigma_x), generated by first sampling from Z
        """
        z = np.random.normal(0, 1, size=(n_samples, self.n_z))
        eps = np.random.normal(0, 1, size=(n_samples, self.n_x))
        data = np.dot(z, self.W.T) + self.b + np.dot(eps, self.chol_sigma_x_z)
        return data 


    def compute_posterior(self, x, z):
        """ Compute the posterior and its parameters: p(z|x), mu_z_x, sigma_z_x
        """
        solver = np.linalg.solve(self.sigma_x, x - self.b)
        mu_z_x = self.mu_z + np.dot(self.W.T, solver)
        sigma_z_x = self.sigma_z - np.dot(self.W.T, np.dot(np.linalg.inv(self.sigma_x), self.W))
        pdf = multivariate_normal.pdf(z, mu_z_x, sigma_z_x)
        return pdf, mu_z_x, sigma_z_x
 

    def compute_log_likelihood(self, x):
        """ True log likelihood of the data
        """
        p_x = multivariate_normal.pdf(x, self.b, self.sigma_x)
        return np.mean(np.log(p_x))


def load(n_x, n_z, n_samples_train, n_samples_test):
    data = {}
    multivarGaussData = MultivariateGaussianData(n_x, n_z)
    generatedX = multivarGaussData.create_dataset(n_samples_train + n_samples_test)
    data["gaussian_class"] = multivarGaussData
    data["train_set_x"] = generatedX[:n_samples_train]
    data["test_set_x"] = generatedX[n_samples_train:]
    data["n_samples_train"] = n_samples_train
    data["n_samples_test"] = n_samples_test
    data["input_size"] = n_x
    data["binary"] = 0
    return data


if __name__=='__main__':
    opts = {"n_x": 10,
            "n_z": 2,
            "n_train": 500,
            "n_test" : 50,
            "n_samples": 5,
            "n_samples_is": 10,
            "n_epochs": 100,
            "period" : 10,
            "alpha": 0.9,
            "learning_rate": 0.0001}

    data = load(opts["n_x"], opts["n_z"], opts["n_train"], opts["n_test"])
    # pdb.set_trace()
