# -*- coding: utf-8 -*-
"""
Created on Wed Jan 12 12:30:23 2022

@author: Rene
"""

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.stats import multivariate_normal

def lims(mplotlims):
    scale = 1.021
    offset = (mplotlims[1] - mplotlims[0])*scale
    return mplotlims[1] - offset, mplotlims[0] + offset

def plot_cube_edges(figure,xlims,ylims,zlims,the_color,the_line_width,the_line_style):
    figure.plot([xlims[1],xlims[1]], [ylims[0],ylims[1]], [zlims[0],zlims[0]],color = the_color,lw =the_line_width,ls = the_line_style)
    figure.plot([xlims[1],xlims[1]], [ylims[0],ylims[1]], [zlims[1],zlims[1]],color = the_color,lw =the_line_width,ls = the_line_style)
    figure.plot([xlims[0],xlims[0]], [ylims[0],ylims[1]], [zlims[0],zlims[0]],color = the_color,lw =the_line_width,ls = the_line_style)
    figure.plot([xlims[0],xlims[0]], [ylims[0],ylims[1]], [zlims[1],zlims[1]],color = the_color,lw =the_line_width,ls = the_line_style)


    figure.plot([xlims[0],xlims[1]], [ylims[1],ylims[1]], [zlims[0],zlims[0]],color = the_color,lw =the_line_width,ls = the_line_style)
    figure.plot([xlims[0],xlims[1]], [ylims[1],ylims[1]], [zlims[1],zlims[1]],color = the_color,lw =the_line_width,ls = the_line_style)
    figure.plot([xlims[0],xlims[1]], [ylims[0],ylims[0]], [zlims[0],zlims[0]],color = the_color,lw =the_line_width,ls = the_line_style)
    figure.plot([xlims[0],xlims[1]], [ylims[0],ylims[0]], [zlims[1],zlims[1]],color = the_color,lw =the_line_width,ls = the_line_style)

    figure.plot([xlims[1],xlims[1]], [ylims[1],ylims[1]], [zlims[0],zlims[1]],color = the_color,lw =the_line_width,ls = the_line_style)
    figure.plot([xlims[1],xlims[1]], [ylims[0],ylims[0]], [zlims[0],zlims[1]],color = the_color,lw =the_line_width,ls = the_line_style)
    figure.plot([xlims[0],xlims[0]], [ylims[1],ylims[1]], [zlims[0],zlims[1]],color = the_color,lw =the_line_width,ls = the_line_style)
    figure.plot([xlims[0],xlims[0]], [ylims[0],ylims[0]], [zlims[0],zlims[1]],color = the_color,lw =the_line_width,ls = the_line_style)

def plot_xy_plane(figure,xlims,ylims,zlims,the_color,the_line_width,the_line_style):
    figure.plot([xlims[1],xlims[1]], [ylims[0],ylims[1]], [0,0],color = the_color,lw =the_line_width,ls = the_line_style)
    figure.plot([xlims[0],xlims[0]], [ylims[0],ylims[1]], [0,0],color = the_color,lw =the_line_width,ls = the_line_style)
    figure.plot([xlims[0],xlims[1]], [ylims[0],ylims[0]], [0,0],color = the_color,lw =the_line_width,ls = the_line_style)
    figure.plot([xlims[0],xlims[1]], [ylims[1],ylims[1]], [0,0],color = the_color,lw =the_line_width,ls = the_line_style)


the_mu = [2,3]
the_sigma = [1,1]
the_rho = 0.75
D = [[the_sigma[0],0],[0,the_sigma[1]]]
pearson_matrix = [[1.0, the_rho], [the_rho, 1]]
cov_matrix = np.matmul(D, np.matmul(pearson_matrix, D))

# Our 2-dimensional distribution will be over variables X and Y
N = 125
X = np.linspace(-3*the_sigma[0]+the_mu[0], 3*the_sigma[0]+the_mu[0], N)

lb_x = 1
ub_x = 3
result_X = np.where((X < lb_x) | (X > ub_x))
result_x = np.where((X > lb_x) & (X < ub_x))

Y = np.linspace(-3*the_sigma[1]+the_mu[1], 3*the_sigma[1]+the_mu[1], N)
lb_y = 2
ub_y = 4
result_Y = np.where((Y < lb_y) | (Y > ub_y))
result_y = np.where((Y > lb_y) & (Y < ub_y))

X, Y = np.meshgrid(X, Y)

# Pack X and Y into a single 3-dimensional array
pos = np.empty(X.shape + (2,))
pos[:, :, 0] = X
pos[:, :, 1] = Y

# Evaluate the multivariate normal pdf
mv_norm = multivariate_normal(the_mu, cov_matrix)
Z = mv_norm.pdf(pos)
max_Z = np.max(Z)

Z_Copy = Z.copy()
Z_Copy[result_X,:] = 0
Z_Copy[:,result_Y] = 0

# Evaluating the marginal distribution of x
x_marg = np.linspace(-3*the_sigma[0]++the_mu[0], 3*the_sigma[0]++the_mu[0], N)
x_pdf = multivariate_normal.pdf(x_marg, mean=the_mu[0], cov=cov_matrix[0,0])
x_mode_pdf = multivariate_normal.pdf(the_mu[0], mean=the_mu[0], cov=cov_matrix[0,0])

# Evaluating the marginal distribution of y
y_marg = np.linspace(-3*the_sigma[1]+the_mu[1], 3*the_sigma[1]++the_mu[1], N)
y_pdf = multivariate_normal.pdf(y_marg, mean=the_mu[1], cov=cov_matrix[1,1])
y_mode_pdf = multivariate_normal.pdf(the_mu[1], mean=the_mu[1], cov=cov_matrix[1,1])


# Adding the left panel with the theoretical pmf
plt.rc('font', family='serif', size='8')
fig = plt.figure()
fig.set_figwidth(9)
fig.set_figheight(5)
fig.subplots_adjust(left=0.075, bottom=0, right=0.975, top=1, hspace=0, wspace=0.2)


ax = fig.add_subplot(121, projection='3d')

# Setting the axes limita
ax.set_zlim(-max_Z,max_Z)
ax.set_xlim(-3*the_sigma[0]+the_mu[0],3*the_sigma[0]+the_mu[0])
ax.set_ylim(-3*the_sigma[1]+the_mu[1],3*the_sigma[1]+the_mu[1])

# Setting the axes labels
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('Probability Density Function')

# Get rid of colored axes planes
ax.xaxis.pane.fill = False
ax.yaxis.pane.fill = False
ax.zaxis.pane.fill = False

# Plot the edges
the_xlims, the_ylims, the_zlims = lims(ax.get_xlim()), lims(ax.get_ylim()), lims(ax.get_zlim())
plot_cube_edges(ax,the_xlims, the_ylims, the_zlims,'darkgray',1,":")

# Plot the xy-plane
plot_xy_plane(ax,the_xlims, the_ylims, the_zlims,'darkgray',1,"-")

# Set the rotation of the figure
ax.view_init(elev=15, azim=35+90+90)

# Bonus: To get rid of the grid as well:
ax.grid(False)

# Plotting the surface with wireframe:
ax.plot_wireframe(X, Y, Z, rstride=3, cstride=3, alpha=0.8,lw=0.5,colors='black')
ax.plot_surface(X, Y, Z, rstride=3, cstride=3, cmap=plt.cm.jet,alpha=0.4)

# Plotting a set number of contours:
n_contours = 15
levels = np.arange(0,max_Z,max_Z/n_contours)
ax.contour(X, Y, Z, zdir='z', offset=the_zlims[0], colors ='black',levels=levels,linewidths=0.5)

# Ploting the marginal distributions of x and y
max_marg = np.max([x_pdf,y_pdf])
ratio = max_Z/max_marg
ax.plot(x_marg, ratio*x_pdf,zs=the_ylims[1], zdir='y', label='curve in (x, z)',color = 'indianred')
ax.plot(y_marg, ratio*y_pdf,zs=the_xlims[1], zdir='x', label='curve in (y, z)',color = 'indianred')

# Plot some line at the mean values
ax.plot([the_mu[0],the_mu[0]], [the_ylims[0],the_ylims[1]], [the_zlims[0],the_zlims[0]],color = 'darkgray',lw =1,ls = ':')
ax.plot([the_xlims[0],the_xlims[1]],[the_mu[1],the_mu[1]], [the_zlims[0],the_zlims[0]],color = 'darkgray',lw =1,ls = ':')

ax.plot([the_mu[0],the_mu[0]], [the_ylims[1],the_ylims[1]], [the_zlims[0],ratio*x_mode_pdf],color = 'darkgray',lw =1,ls = ':')
ax.plot([the_xlims[1],the_xlims[1]],[the_mu[1],the_mu[1]], [the_zlims[0],ratio*y_mode_pdf],color = 'darkgray',lw =1,ls = ':')

# adding a title to the plot
text_str = 'Bivariate Normal pdf'
ax.text2D(0.2, 0.9, text_str, transform=ax.transAxes,size=15)

# adding panel letter to the plot
text_str = 'A'
ax.text2D(0, 0.9, text_str, transform=ax.transAxes,size=15,color = 'red')

#***** Adding the right subplot

ax = fig.add_subplot(122, projection='3d')

# Setting the axes limita
ax.set_zlim(-max_Z,max_Z)
ax.set_xlim(-3*the_sigma[0]+the_mu[0],3*the_sigma[0]+the_mu[0])
ax.set_ylim(-3*the_sigma[1]+the_mu[1],3*the_sigma[1]+the_mu[1])

# Setting the axes labels
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('Probability Density Function')

# Get rid of colored axes planes
ax.xaxis.pane.fill = False
ax.yaxis.pane.fill = False
ax.zaxis.pane.fill = False

# Plot the edges
the_xlims, the_ylims, the_zlims = lims(ax.get_xlim()), lims(ax.get_ylim()), lims(ax.get_zlim())
plot_cube_edges(ax,the_xlims, the_ylims, the_zlims,'darkgray',1,":")

# Plot the xy-plane
plot_xy_plane(ax,the_xlims, the_ylims, the_zlims,'darkgray',1,"-")

# Set the rotation of the figure
ax.view_init(elev=15, azim=35+90+90)

# Bonus: To get rid of the grid as well:
ax.grid(False)

# Plotting the surface with wireframe:
ax.plot_wireframe(X, Y, Z_Copy, rstride=3, cstride=3, alpha=0.8,lw=0.5,colors='black')
ax.plot_surface(X, Y, Z_Copy, rstride=3, cstride=3, cmap=plt.cm.jet,alpha=0.4)

# Plotting a set number of contours:
n_contours = 15
levels = np.arange(0,max_Z,max_Z/n_contours)
ax.contour(X, Y, Z_Copy, zdir='z', offset=the_zlims[0], colors ='black',levels=levels,linewidths=0.5)

# Ploting the marginal distributions of x and y
max_marg = np.max([x_pdf,y_pdf])
ratio = max_Z/max_marg
ax.plot(x_marg[result_x], ratio*x_pdf[result_x],zs=the_ylims[1], zdir='y', label='curve in (x, z)',color = 'indianred')
ax.plot(y_marg[result_y], ratio*y_pdf[result_y],zs=the_xlims[1], zdir='x', label='curve in (y, z)',color = 'indianred')

# Plot some line at the mean values
ax.plot([the_mu[0],the_mu[0]], [the_ylims[0],the_ylims[1]], [the_zlims[0],the_zlims[0]],color = 'darkgray',lw =1,ls = ':')
ax.plot([the_xlims[0],the_xlims[1]],[the_mu[1],the_mu[1]], [the_zlims[0],the_zlims[0]],color = 'darkgray',lw =1,ls = ':')

ax.plot([the_mu[0],the_mu[0]], [the_ylims[1],the_ylims[1]], [the_zlims[0],ratio*x_mode_pdf],color = 'darkgray',lw =1,ls = ':')
ax.plot([the_xlims[1],the_xlims[1]],[the_mu[1],the_mu[1]], [the_zlims[0],ratio*y_mode_pdf],color = 'darkgray',lw =1,ls = ':')

# Plot some line at the lb_x and ub_x
ax.plot([lb_x,lb_x], [the_ylims[0],the_ylims[1]], [the_zlims[0],the_zlims[0]],color = 'darkgray',lw =1,ls = '--')
ax.plot([ub_x,ub_x], [the_ylims[0],the_ylims[1]], [the_zlims[0],the_zlims[0]],color = 'darkgray',lw =1,ls = '--')

pdf_value = multivariate_normal.pdf(lb_x, mean=the_mu[0], cov=cov_matrix[0,0])
ax.plot([lb_x,lb_x], [the_ylims[1],the_ylims[1]], [the_zlims[0],ratio*pdf_value],color = 'darkgray',lw =1,ls = '--')
pdf_value = multivariate_normal.pdf(ub_x, mean=the_mu[0], cov=cov_matrix[0,0])
ax.plot([ub_x,ub_x], [the_ylims[1],the_ylims[1]], [the_zlims[0],ratio*pdf_value],color = 'darkgray',lw =1,ls = '--')

# Plot some line at the lb_y and ub_y
ax.plot([the_xlims[0],the_xlims[1]],[lb_y,lb_y], [the_zlims[0],the_zlims[0]],color = 'darkgray',lw =1,ls = '--')
ax.plot([the_xlims[0],the_xlims[1]],[ub_y,ub_y], [the_zlims[0],the_zlims[0]],color = 'darkgray',lw =1,ls = '--')


pdf_value = multivariate_normal.pdf(lb_y, mean=the_mu[1], cov=cov_matrix[1,1])
ax.plot([the_xlims[1],the_xlims[1]],[lb_y,lb_y], [the_zlims[0],ratio*pdf_value],color = 'darkgray',lw =1,ls = '--')
pdf_value = multivariate_normal.pdf(ub_y, mean=the_mu[1], cov=cov_matrix[1,1])
ax.plot([the_xlims[1],the_xlims[1]],[ub_y,ub_y], [the_zlims[0],ratio*pdf_value],color = 'darkgray',lw =1,ls = '--')


# adding a title to the plot
text_str = 'Bivariate Normal pdf'
ax.text2D(0.2, 0.9, text_str, transform=ax.transAxes,size=15)

# adding panel letter to the plot
text_str = 'B'
ax.text2D(0, 0.9, text_str, transform=ax.transAxes,size=15,color = 'red')

bbox = fig.bbox_inches.from_bounds(0, 0.5, 9, 4)

plt.savefig("Bivariate_Normal_pdf.png",  bbox_inches=bbox, dpi=300)