# -*- 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

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)

def f_x_marg_pdf(x):
    value=0
    if ((x>=0)&(x<=3)):
        value=(2/225)*(9*x*x+7*x) 
    return(value)

def f_y_marg_pdf(y):
    value=0
    if ((y>=1)&(y<=2)):
        value=(9/75)*(y*y+4*y)
    return(value) 

def f_joint_pdf(x,y):
    value=0
    if ((x>=0)&(x<=3)&(y>=1)&(y<=2)):
        value=(2/75)*(2*x*x*y+x*y*y)
    return(value) 


# Our 2-dimensional distribution will be over variables X and Y
N = 50

lb_x = -1
ub_x = 4
X = np.linspace(lb_x, ub_x, N)

lb_y = 0
ub_y = 3
Y = np.linspace(lb_y, ub_y, N)

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

Z = X.copy()

# Evaluate the multivariate pdf
for i in np.arange(0,N):
    for j in np.arange(0,N):
        x_value = X[i,j]
        y_value = Y[i,j]
        Z[i,j] = f_joint_pdf(x_value,y_value)

max_Z = np.max(Z)

# Evaluating the marginal distribution of x
x_marg = np.linspace(0, 3, N)
x_pdf = x_marg.copy()
for i in np.arange(0,N):
    x_pdf[i] = f_x_marg_pdf(x_marg[i])

# Evaluating the marginal distribution of y
y_marg = np.linspace(1, 2, N)
y_pdf = y_marg.copy()
for j in np.arange(0,N):
    y_pdf[j] = f_y_marg_pdf(y_marg[j])

# ratio for plotting the marginal distributions in the facets
max_marg = np.max([x_pdf,y_pdf])
ratio = max_Z/max_marg

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

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

# Setting the axes limita
ax.set_zlim(-max_Z,max_Z)
ax.set_xlim(lb_x,ub_x)
ax.set_ylim(lb_y,ub_y)

# 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=1, cstride=1, alpha=0.8,lw=0.5,colors='black')
ax.plot_surface(X, Y, Z, rstride=1, cstride=1, 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
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 lines at the boundaries
ax.plot([0,0], [the_ylims[0],the_ylims[1]], [the_zlims[0],the_zlims[0]],color = 'darkgray',lw =1,ls = '--')
ax.plot([3,3], [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]],[1,1], [the_zlims[0],the_zlims[0]],color = 'darkgray',lw =1,ls = '--')
ax.plot([the_xlims[0],the_xlims[1]],[2,2], [the_zlims[0],the_zlims[0]],color = 'darkgray',lw =1,ls = '--')

pdf_value = f_x_marg_pdf(0) 
ax.plot([0,0], [the_ylims[1],the_ylims[1]], [the_zlims[0],ratio*pdf_value],color = 'darkgray',lw =1,ls = '--')
pdf_value = f_x_marg_pdf(3) 
ax.plot([3,3], [the_ylims[1],the_ylims[1]], [the_zlims[0],ratio*pdf_value],color = 'darkgray',lw =1,ls = '--')

pdf_value = f_y_marg_pdf(1) 
ax.plot([the_xlims[1],the_xlims[1]],[1,1], [the_zlims[0],ratio*pdf_value],color = 'darkgray',lw =1,ls = '--')
pdf_value = f_y_marg_pdf(2) 
ax.plot([the_xlims[1],the_xlims[1]],[2,2], [the_zlims[0],ratio*pdf_value],color = 'darkgray',lw =1,ls = '--')

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

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

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