'''
Created on 10 ago 2017

@author: davide
'''
from astropy.stats.tests import test_sigma_clipping

'''
    Compute and plot an example of bivariate gaussian distribution
'''

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import axes3d

N = 60 # number of samples
X = np.linspace(-3, 3, N) 
Y = np.linspace(-3, 4, N)
X, Y = np.meshgrid(X, Y)

# Mean vector and Covariance matrix
mu = np.array([0., 1.])
sigma = np.array([[1., -0.5], [-0.5, 1.5]]) # Should be symmetric and positive definite...

pos = np.empty(X.shape + (2, ))
pos[:, :, 0] = X
pos[:, :, 1] = Y

def multivariate_gaussian(pos, mu, sigma):
    '''
        Return the multivariate gaussian for the position in array pos
    '''
    
    n = mu.shape[0]
    sigma_det = np.linalg.det(sigma)
    sigma_inv = np.linalg.inv(sigma)
    scale_factor = np.sqrt((2*np.pi)**n * sigma_det)
    scale_factor = 1/scale_factor
    fac = np.einsum('...k,kl,...l->...', pos - mu, sigma_inv, pos - mu)
    return np.exp(-fac / 2) * scale_factor
    
Z = multivariate_gaussian(pos, mu, sigma)

fig = plt.figure()
ax = fig.gca(projection='3d')
ax.plot_surface(X, Y, Z, rstride = 3, cstride = 3, linewidth = 1, antialiased=True, cmap = cm.viridis)
cset = ax.contourf(X, Y, Z, zdir = 'z', offset = -0.15, cmap = cm.viridis)

# Adjust the limits, ticks and view angle
ax.set_zlim(-0.15,0.2)
ax.set_zticks(np.linspace(0,0.2,5))
ax.view_init(27, -21)

plt.show()

Z1 = Z

'''
    Now use scipy.stats 
'''
from scipy.stats import multivariate_normal
F = multivariate_normal(mu, sigma)
Z = F.pdf(pos)
fig = plt.figure()
ax = fig.gca(projection='3d')
ax.plot_surface(X, Y, Z, rstride = 3, cstride = 3, linewidth = 1, antialiased=True, cmap = cm.viridis)
cset = ax.contourf(X, Y, Z, zdir = 'z', offset = -0.15, cmap = cm.viridis)

# Adjust the limits, ticks and view angle
ax.set_zlim(-0.15,0.2)
ax.set_zticks(np.linspace(0,0.2,5))
ax.view_init(27, -21)

plt.show()

'''
    Compare the two ways of computing Z
'''
delta = Z1-Z 
print(delta)
print()
print("L2 norm of the difference matrix: " + str(np.linalg.norm(delta)) + " should be (almost) zero...")

'''
    Now the other way around:
    We have the data, assume they are generated by a set of iid multivariate gaussians
    then estimate the parameters of these
    We're gonna use the ML estimator. for Gaussian distributions it is known in closed form
'''

# data points
N = int(1E6)
X = np.random.multivariate_normal(mu, sigma, N)
print(X.shape)
print(X)

'''
    np.random.multivariate_normal is N x 2,
    i. e. each sample is in the rows 
'''
mu_hat = np.mean(X, axis = 0)
print()
print("Estimated mean:")
print(mu_hat.shape)
print(mu)
print(mu_hat)
delta_mu = mu - mu_hat
print("Actual mu: " + str(mu) + ", estimated (ML) mu_hat: " + str(mu_hat) + ", L2 norm of the difference: " + str(np.linalg.norm(delta_mu)) + ".")

'''
    Estimating the covariance matrix
'''
DEBUG = 0

d = X.shape[1]
sigma_hat = np.zeros(shape = (d, d))
for i in range(N):
    temp = np.matrix(X[i, :] - mu_hat)
    temp_T = np.transpose(temp)
    mult = temp_T.dot(temp)
    sigma_hat = sigma_hat + mult
    
    if (DEBUG > 0):
        print("temp: " + str(temp))
        print(temp.shape)
        
        print("temp_T: " + str(temp_T))
        print(temp_T.shape)
        
        print("mult: " + str(mult))
        print(mult.shape)
        
sigma_hat = sigma_hat / N

print()
print("Estimated covariance:")
print(sigma)
print(sigma_hat)
delta_sigma = sigma - sigma_hat
print("Actual sigma: " + str(sigma) + ", estimated (ML) sigma_hat: " + str(sigma_hat) + ", L2 norm of the difference: " + str(np.linalg.norm(delta_sigma)) + ".")

'''
    Now use library functions to estimate covariance matrix
'''
mu_hat = np.mean(X, axis = 0)

d = X.shape[1]
sigma_hat = np.eye(d)
Phi2 = multivariate_normal.pdf(X, mean = mu_hat, cov = sigma_hat)
print()

print('Phi2.shape: ' + str(Phi2.shape))
print(Phi2)