import numpy as np
import math
import scipy.stats as stats
from abc import ABCMeta, abstractmethod
import distributions 
import utils_math
from problems import ABC_problems
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import time


class Gaussian_Copula_Problem(ABC_problems.ABC_Problem):

    '''
    A problem where data is generated by a 2D Gaussian copula where: marginal 1 = beta, marginal 2 = MoG
    '''

    def __init__(self, N=100, n=50):
        self.N = N                                                                       # number of posterior samples
        self.n = n                                                                       # number of i.i.d data x_i ~ p(x|theta)

        self.prior = [distributions.uniform, distributions.uniform, distributions.uniform]
        self.prior_args = np.array([[0.5, 12.5], [0, 1.0], [0.4, 0.8]])                  # uniform prior
        self.simulator_args = ['alpha', 'coeff', 'cov']                                  # just for information
        self.K = 3                                                                       # number of parameters
        self.stat = 'raw'
        
        self.true_cov = np.array([[1, 0.60], [0.60, 1]])
        self.true_coeff = 0.50
        self.true_alpha = 6.0
        self.true_beta = 2.0
        self.true_mean1 = 1.0
        self.true_mean2 = 4.0
        self.true_sigma1 = 1.0
        self.true_sigma2 = 0.5

        self.MoG_u2x_mappings = distributions.MoG2.learn_u2x_mappings(self.true_mean1, self.true_sigma1, self.true_mean2, self.true_sigma2)

    def get_true_theta(self):
        return np.array([self.true_alpha, self.true_coeff, self.true_cov[0, 1]])

    
    # A. (normalized) marginal quantiles
    def _ss_quantiles(self, X, n_quantiles):
        dim = X.shape[1]
        prob = np.linspace(0.025, 0.975, n_quantiles)
        stat = np.zeros([1, n_quantiles*dim])
        for k in range(dim):
            quantiles = stats.mstats.mquantiles(X[:, k], prob)
            stat_k = quantiles
            stat[0, k*n_quantiles : (k+1)*n_quantiles] = np.array(stat_k)
        return stat

    # B. correlation between latent
    def _ss_corr(self, Z):
        V = np.mat(Z).T * np.mat(Z) / Z.shape[0]
        (d,d) = V.shape
        upper_tri_elements = V[np.triu_indices(d, k=1)]
        stat = np.array(upper_tri_elements)
        return stat
    
    def statistics(self, data, theta=None):
        if self.stat == 'raw':
            # (marginal quantiles) + (latent correlation) as summary statistics
            stat_A = self._ss_quantiles(data, n_quantiles=20)
            stat_B = self._ss_corr(self.X2Z(data, theta[0], theta[1]))
            stat = np.hstack((stat_A, stat_B))
            return stat
        else:
            # (marginal quantiles) + (latent correlation) as summary statistics
            stat_A = self._ss_quantiles(data, n_quantiles=5)
            stat_B = self._ss_corr(self.X2Z(data, theta[0], theta[1]))
            stat = np.hstack((stat_A, stat_B))
            return stat
            
    def simulator(self, theta):
        # some preparation
        alpha = theta[0]
        coeff = theta[1]
        V = np.array([[1, theta[2]], [theta[2], 1]])

        # sample z ~ N(0, V)
        Z = distributions.normal_nd.draw_samples([0, 0], V, self.n)

        # convert z to x
        X = self.Z2X(Z, alpha, coeff)
        return X

    def log_likelihood(self, theta):

        # calculate L(theta; x_o) = p(theta|x_o)
        
        alpha = theta[0]
        coeff = theta[1]
        V = np.array([[1, theta[2]], [theta[2], 1]])

        # compute the copula density
        Z = self.X2Z(self.data_obs, alpha, coeff)
        c = distributions.copula.copula_density(Z, V)

        # compute the marginal pdf
        p1 = distributions.beta.pdf(self.data_obs[:, 0], alpha, self.true_beta)
        p2 = distributions.MoG2.pdf(self.data_obs[:, 1], coeff, self.true_mean1, self.true_sigma1, self.true_mean2, self.true_sigma2)

        # likelihood = copula density * marginal pdf
        ll = (np.log(c) + np.log(p1) + np.log(p2)).sum()
        return ll
    
    def log_pdf(self, data, theta):

        # calculate p(x|theta) 

        alpha = theta[0]
        coeff = theta[1]
        V = np.array([[1, theta[2]], [theta[2], 1]])

        # compute the copula density
        Z = self.X2Z(data, alpha, coeff)
        c = distributions.copula.copula_density(Z, V)

        # compute the marginal pdf
        p1 = distributions.beta.pdf(data[:, 0], alpha, self.true_beta)
        p2 = distributions.MoG2.pdf(data[:, 1], coeff, self.true_mean1, self.true_sigma1, self.true_mean2, self.true_sigma2)

        # copula density * marginal pdf = pdf
        return np.log(c)+np.log(p1)+np.log(p2)

    def sample_from_prior(self):
        sample_alpha = self.prior[0].draw_samples(self.prior_args[0, 0], self.prior_args[0, 1],  1)[0]
        sample_coeff = self.prior[1].draw_samples(self.prior_args[1, 0], self.prior_args[1, 1],  1)[0]
        sample_cov = self.prior[2].draw_samples(self.prior_args[2, 0], self.prior_args[2, 1],  1)[0]
        return np.array([sample_alpha, sample_coeff, sample_cov])

    def X2Z(self, X, alpha, coeff):
        # get u = beta-CDF(x), MoG-CDF(x)
        U = np.zeros(X.shape)
        U[:, 0] = distributions.beta.cdf(X[:, 0], alpha, self.true_beta)
        U[:, 1] = distributions.MoG2.cdf(X[:, 1], coeff, self.true_mean1, self.true_sigma1, self.true_mean2, self.true_sigma2)

        # get z = inverse-Phi(u)
        Z = distributions.normal.invcdf(U, mu=0, sigma=1)
        return Z

    def Z2X(self, Z, alpha, coeff):
        # get u = Phi(z)
        U = np.zeros(Z.shape)
        X = np.zeros(Z.shape)
        for k in range(2):
            U[:, k] = distributions.normal.cdf(Z[:, k], 0, 1)

        # get x = inverse-CDF(u)
        X[:, 0] = distributions.beta.invcdf(U[:, 0], alpha, self.true_beta)
        X[:, 1] = distributions.MoG2.invcdf(U[:, 1], coeff, self.MoG_u2x_mappings)
        return X

    def visualize(self):
        print('visualizing p(x|theta)')
        
        # preparation
        samples = self.data_obs
        [m, dim] = samples.shape
        min_values = samples.min(axis=0)
        max_values = samples.max(axis=0)

        # likelihood values 
        N_grid = 300
        ranges = []
        for k in range(dim):
            r = np.array(np.linspace(min_values[k], max_values[k], N_grid))
            ranges.append(r)
        X, Y = np.meshgrid(*ranges)
        R = np.array(np.meshgrid(*ranges)).T.reshape(-1, dim)

        pdf = np.exp(self.log_pdf(R, self.get_true_theta()))
        Z = pdf.reshape(X.shape)

        # plot the contour
        plt.figure(figsize=(5, 5))
        plt.contour(X, Y, Z, 10, cmap='jet', linewidths=0.75)
        plt.xlabel(r'$y_1$')
        plt.ylabel(r'$y_2$')
        plt.xlim((0.3, 1.1))
        plt.ylim((-1, 6))
        plt.show()
        