# -*- coding: utf-8 -*-
"""
Created on Fri Jan 14 14:49:00 2022

@author: Rene
"""

import numpy as np
import pandas as pd
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from matplotlib.cm import ScalarMappable
import Dist_Library as dl

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)


# S is the sum of the two die outcomes

# M is the maximum of the two die outcomes

# creating a dataframe with all possible combintations of Die_1, Die_2, S and M
Outcome_Die_1 = np.zeros(36)
Outcome_Die_2 = np.zeros(36)
Outcome_M = np.zeros(36)
Outcome_S = np.zeros(36)
k= -1
for i in np.arange(0,6):
  for j in np.arange(0,6):
    k = k+1
    Outcome_Die_1[k] = i+1
    Outcome_Die_2[k] = j+1
    Outcome_S[k] = i+j+2
    Outcome_M[k] = max(i+1,j+1)
    
 
Outcome_S_M = pd.DataFrame(data = Outcome_Die_1) 
Outcome_S_M.columns = ['Die 1']   
Outcome_S_M['Die 2'] = Outcome_Die_2
Outcome_S_M['S'] = Outcome_S
Outcome_S_M['M'] = Outcome_M

# Initializing the Joint Probability Mass Function of S and M to value zero
Joint_PMF_S_M=np.zeros((11,6))
S_value = np.zeros(11)
M_value = np.zeros(6)
for i in np.arange(2,13):
  for j in np.arange(1,7):
      S_value[i-2] = i
      M_value[j-1] = j
      Query_Str = 'S == ' + str(i) +' & ' + 'M == ' +str(j)
      Outcomes = Outcome_S_M.query(Query_Str)
      n_outcomes = len(Outcomes)
      Joint_PMF_S_M[i-2,j-1] = n_outcomes/36

#Evaluating thet Joint CDF and the marginal distributions
Joint_CDF_S_M = np.zeros((11,6))
Marginal_PMF_M = np.zeros(6)
Marginal_PMF_S = np.zeros(11)
for i in np.arange(0,11):
    for j in np.arange(0,6):
        Joint_CDF_S_M[i,j] = Joint_PMF_S_M[:i+1,:j+1].sum()
        Marginal_PMF_S[i] = Joint_PMF_S_M[i,:6].sum()
        Marginal_PMF_M[j] = Joint_PMF_S_M[:11,j].sum()


plt.rc('font', family='serif', size='8')
fig = plt.figure(0)
fig.set_figwidth(9)
fig.set_figheight(5)
fig.subplots_adjust(left=0.1, bottom=0.1, right=0.9, top=0.9, hspace=0, wspace=0.2)


# Plotting the Joint Probability Mass Function
data_2d = Joint_PMF_S_M
#
# Convert it into an numpy array.
#
data_array = np.array(data_2d)
#
# Create a figure for plotting the data as a 3D histogram.
#
ax = fig.add_subplot(121, projection='3d')

# Setting the axes limita
ax.set_zlim(0,np.max(data_2d))
ax.set_xlim(0.5,6.5)
ax.set_ylim(1.5,12.5)

# Setting the axes labels
ax.set_xlabel('M')
ax.set_ylabel('S')
ax.set_zlabel('Probability Mass Function')

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

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

# Set the rotation of the figure
ax.view_init(elev=35, azim=90+90+20)
#
# Create an X-Y mesh of the same dimension as the 2D data. You can
# think of this as the floor of the plot.
#
x_data, y_data = np.meshgrid( np.arange(data_array.shape[1]),
                              np.arange(data_array.shape[0]) )
#
# Flatten out the arrays so that they may be passed to "ax.bar3d".
# Basically, ax.bar3d expects three one-dimensional arrays:
# x_data, y_data, z_data. The following call boils down to picking
# one entry from each array and plotting a bar to from
# (x_data[i], y_data[i], 0) to (x_data[i], y_data[i], z_data[i]).
#

x_data = x_data.flatten()+0.5
y_data = y_data.flatten()+1.5
z_data = data_array.flatten()


max_height = np.max(z_data)   # get range of colorbars so we can normalize
min_height = np.min(z_data)

# scale each z to [0,1], and get their rgb values
cmap = plt.cm.get_cmap('jet') # Get desired colormap - you can change this!
rgba = [cmap((k-min_height)/max_height) for k in z_data]

# define the bins and normalize
bounds = np.linspace(0, max(z_data), 36)
sm = ScalarMappable(cmap=cmap, norm=plt.Normalize(0,max(z_data)))
sm.set_array([])

v1 = np.linspace(z_data.min(), z_data.max(), 8, endpoint=True)      

# Plot the Color Legend
cbar = plt.colorbar(sm, boundaries=bounds,ticks=v1,shrink=0.75)
cbar.ax.set_yticklabels(["{:4.2f}".format(i) for i in v1]) # add the labels

# Plot the 3d Historgram
ax.bar3d( x_data,
          y_data,
          np.zeros(len(z_data)),
          1, 1, z_data, color=rgba, edgecolor = 'black' )

# 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,":")

# adding a title to the plot
text_str = 'Joint PMF'
ax.text2D(0.25, 1, text_str, transform=ax.transAxes,size=15)


# Plotting the Joint Probability Cumulative Distribution Function
data_2d = Joint_CDF_S_M
#
# Convert it into an numpy array.
#
data_array = np.array(data_2d)
#
# Create a figure for plotting the data as a 3D histogram.
#

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

# Setting the axes limita
ax.set_zlim(0,np.max(data_2d))
ax.set_xlim(0.5,6.5)
ax.set_ylim(1.5,12.5)

# Setting the axes labels
ax.set_xlabel('M')
ax.set_ylabel('S')
ax.set_zlabel('Probability Mass Function')

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

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

# Set the rotation of the figure
ax.view_init(elev=35, azim=90+90+20)
#
# Create an X-Y mesh of the same dimension as the 2D data. You can
# think of this as the floor of the plot.
#
x_data, y_data = np.meshgrid( np.arange(data_array.shape[1]),
                              np.arange(data_array.shape[0]) )
#
# Flatten out the arrays so that they may be passed to "ax.bar3d".
# Basically, ax.bar3d expects three one-dimensional arrays:
# x_data, y_data, z_data. The following call boils down to picking
# one entry from each array and plotting a bar to from
# (x_data[i], y_data[i], 0) to (x_data[i], y_data[i], z_data[i]).
#

x_data = x_data.flatten()+0.5
y_data = y_data.flatten()+1.5
z_data = data_array.flatten()


max_height = np.max(z_data)   # get range of colorbars so we can normalize
min_height = np.min(z_data)

# scale each z to [0,1], and get their rgb values
cmap = plt.cm.get_cmap('jet') # Get desired colormap - you can change this!
rgba = [cmap((k-min_height)/max_height) for k in z_data]

# define the bins and normalize
bounds = np.linspace(0, max(z_data), 36)
sm = ScalarMappable(cmap=cmap, norm=plt.Normalize(0,max(z_data)))
sm.set_array([])

v1 = np.linspace(z_data.min(), z_data.max(), 8, endpoint=True)      

# Plot the Color Legend
cbar = plt.colorbar(sm, boundaries=bounds,ticks=v1,shrink=0.75)
cbar.ax.set_yticklabels(["{:4.2f}".format(i) for i in v1]) # add the labels

# Plot the 3d Historgram
ax.bar3d( x_data,
          y_data,
          np.zeros(len(z_data)),
          1, 1, z_data, color=rgba, edgecolor = 'black' )

# 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,":")

# adding a title to the plot
text_str = 'Joint CDF'
ax.text2D(0.25, 1, text_str, transform=ax.transAxes,size=15)

bbox = fig.bbox_inches.from_bounds(0.25, 0.75, 8, 3.5)

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


# Plotting the marginal distributions of S

# Creating the probability mass function figure
plt.rc('font', family='serif', size='10')
S_figure = plt.figure(1)
S_figure.set_figwidth(7)
S_figure.set_figheight(3.5)
S_figure.subplots_adjust(left=0.15, bottom=0.15, right=0.95, top=0.9, hspace=0.1, wspace=0.3)

# Adding the left panel with the theoretical pmf
Left_Panel = S_figure.add_subplot(1,2,1)
plt.scatter(S_value,Marginal_PMF_S,color='red')
plt.xlim((1,13))
plt.ylim((-0.01,0.2))
plt.xlabel('s')
plt.ylabel('Probability Mass Function $p(s)$')
plt.title('Marginal PMF - S')
plt.axvline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.axhline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.vlines(S_value, 0, Marginal_PMF_S, color='red', linestyles='-', lw=2)
plt.xticks([2, 4, 6, 8, 10, 12])

# Adding the right panel with the empirical pmf
Right_Panel = S_figure.add_subplot(1,2,2)
dl.plot_discrete_cdf(Right_Panel,S_value,Marginal_PMF_S,'red',2)
plt.xlim((1,13))
plt.ylim((-0.05,1.1))
plt.xlabel('s')
plt.ylabel('Cumulative Distr. Function $F(s) = Pr(S \leq s)$')
plt.title('Marginal CDF - S')
plt.axvline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.axhline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.xticks([2, 4, 6, 8, 10, 12])

plt.savefig('S_PMF_CDF.png', dpi=300)

# Plotting the marginal distributions of M

# Creating the probability mass function figure
plt.rc('font', family='serif', size='10')
M_figure = plt.figure(2)
M_figure.set_figwidth(7)
M_figure.set_figheight(3.5)
M_figure.subplots_adjust(left=0.15, bottom=0.15, right=0.95, top=0.9, hspace=0.1, wspace=0.3)

# Adding the left panel with the theoretical pmf
Left_Panel = M_figure.add_subplot(1,2,1)
plt.scatter(M_value,Marginal_PMF_M,color='steelblue')
plt.xlim((0,7))
plt.ylim((-0.0125,0.4))
plt.xlabel('m')
plt.ylabel('Probability Mass Function $p(s)$')
plt.title('Marginal PMF - M')
plt.axvline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.axhline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.vlines(M_value, 0, Marginal_PMF_M, color='steelblue', linestyles='-', lw=2)
plt.xticks([1, 2, 3, 4, 5, 6])

# Adding the right panel with the empirical pmf
Right_Panel = M_figure.add_subplot(1,2,2)
dl.plot_discrete_cdf(Right_Panel,M_value,Marginal_PMF_M,'steelblue',2)
plt.xlim((0,7))
plt.ylim((-0.05,1.1))
plt.xlabel('s')
plt.ylabel('Cumulative Distr. Function $F(s) = Pr(S \leq s)$')
plt.title('Marginal CDF - M')
plt.axvline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.axhline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.xticks([1, 2, 3, 4, 5, 6])

plt.savefig('M_PMF_CDF.png', dpi=300)