# -*- 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_joint_cdf(x,y):
    value = 0
    if ((x>=0)&(x<=3)&(y>=1)&(y<=2)):
        value1 = (2/75)*((x**3*y**2)/3+(x**2*y**3)/6)
        value2 = (2/75)*((x**3)/3+(x**2)/6)
        value = value1-value2
    if ((x>3)&(y>=1)&(y<=2)):
         value1 = (2/75)*((3**3*y**2)/3+(3**2*y**3)/6)
         value2 = (2/75)*((3**3)/3+(3**2)/6)
         value = value1-value2
    if ((x>=0)&(x<=3)&(y>2)): 
         value1 = (2/75)*((x**3*2**2)/3+(x**2*2**3)/6)
         value2 = (2/75)*((x**3)/3+(x**2)/6)
         value = value1-value2
    if ((x>3)&(y>2)):
        value = 1
    return(value)

def F_x_marg_cdf(x):
    value = F_joint_cdf(x,2)
    return(value)

def F_y_marg_cdf(y):
    value = F_joint_cdf(3,y)
    return(value)


# Our 2-dimensional distribution will be over variables X and Y
# 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)
result_X = np.where((X < 1) | (X > 2))
result_x = np.where((X > 1) & (X < 2))

lb_y = 0
ub_y = 3
Y = np.linspace(lb_y, ub_y, N)
result_Y = np.where((Y < 4/3) | (Y > 5/3))
result_y = np.where((Y > 4/3) & (Y < 5/3))

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_cdf(x_value,y_value)

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(-1, 4, N)
x_cdf = x_marg.copy()
for i in np.arange(0,N):
    x_cdf[i] = F_x_marg_cdf(x_marg[i])
    
# Evaluating the marginal distribution of y
y_marg = np.linspace(0, 3, N)
y_cdf = y_marg.copy()
for j in np.arange(0,N):
    y_cdf[j] = F_y_marg_cdf(y_marg[j])
    
# ratio for plotting the marginal distributions in the facets
max_marg = np.max([x_cdf,y_cdf])
ratio = max_Z/max_marg   

# Adding the left panel with the theoretical cmf
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(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('Cumulative Distribution 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 = 20
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_cdf,zs=the_ylims[1], zdir='y', label='curve in (x, z)',color = 'indianred')
ax.plot(y_marg, ratio*y_cdf,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_cdf(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_cdf(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_cdf(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_cdf(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 cdf'
ax.text2D(0.3, 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(-1,1)
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('Cumulative Distribution 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=1, cstride=1, alpha=0.8,lw=0.5,colors='black')
ax.plot_surface(X, Y, Z_Copy, rstride=1, cstride=1, cmap=plt.cm.jet,alpha=0.4)

# Plotting a set number of contours:
n_contours = 20
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
ax.plot(x_marg[result_x], ratio*x_cdf[result_x],zs=the_ylims[1], zdir='y', label='curve in (x, z)',color = 'indianred')
ax.plot(y_marg[result_y], ratio*y_cdf[result_y],zs=the_xlims[1], zdir='x', label='curve in (y, z)',color = 'indianred')

ax.plot([1,1], [the_ylims[0],the_ylims[1]], [the_zlims[0],the_zlims[0]],color = 'darkgray',lw =1,ls = '--')
ax.plot([2,2], [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]],[4/3,4/3], [the_zlims[0],the_zlims[0]],color = 'darkgray',lw =1,ls = '--')
ax.plot([the_xlims[0],the_xlims[1]],[5/3,5/3], [the_zlims[0],the_zlims[0]],color = 'darkgray',lw =1,ls = '--')

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

pdf_value = F_y_marg_cdf(4/3) 
ax.plot([the_xlims[1],the_xlims[1]],[4/3,4/3], [the_zlims[0],ratio*pdf_value],color = 'darkgray',lw =1,ls = '--')
pdf_value = F_y_marg_cdf(5/3) 
ax.plot([the_xlims[1],the_xlims[1]],[5/3,5/3], [the_zlims[0],ratio*pdf_value],color = 'darkgray',lw =1,ls = '--')


# adding a title to the plot
text_str = 'Joint cdf'
ax.text2D(0.3, 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("Joint_cdf.png",  bbox_inches=bbox, dpi=300)