# -*- coding: utf-8 -*-
"""
Created on Thu Feb  3 12:16:42 2022

@author: Rene
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import geom

def pmf_convolution_geom(x):
    z = np.arange(2,2*len(x))
    Pr_z = np.zeros(len(z))
    Cum_Pr_z = np.zeros(len(z))
    y = np.arange(1,len(x))
    Pr_y = geom.pmf(y, p)
    Cum_Pr = 0
    for i in np.arange(0,len(z)):
        x = z[i]-y
        Pr_x = geom.pmf(x,p)
        Pr_z[i] = sum(Pr_x*Pr_y)
        Cum_Pr = Cum_Pr+Pr_z[i]
        Cum_Pr_z[i] = Cum_Pr
    value = pd.DataFrame({'z':z, 'Pr_z':Pr_z, 'Cum_Pr_z': Cum_Pr_z})
    return(value)

p = 1/75
x = np.arange(1,1000)
pmf_x = geom.pmf(x, p)
cdf_x = geom.cdf(x, p)
pmf_z = pmf_convolution_geom(x)

# Plotting a two panel histogram plot of the speed of light

plt.rc('font', family='serif', size='10')
plt.figure(figsize=(10, 5))
Figure = plt.figure(0)
Figure.set_figwidth(9)
Figure.set_figheight(4.75)
Figure.subplots_adjust(hspace=0.6, wspace=0.4)
off_set_x = 10

# Adding the left panel with the theoretical pmf
Panel = Figure.add_subplot(2,2,1)  
LB_x = 0
UB_x = 500
off_set_x = 0.04*(UB_x-LB_x)
x_lims = (LB_x-off_set_x,UB_x+off_set_x)
LB_y = 0
UB_y = max(pmf_x)
off_set_y = 0.02*(UB_y-LB_y)
y_lims = (LB_y-off_set_y,UB_y+off_set_y)

the_x = x
the_y = pmf_x
plt.xlim(x_lims)
plt.ylim(y_lims)
plt.xlabel('x (in Days)')
plt.ylabel('PMF $p(x)$')
title_str = r'PMF : $X$ or $Y \sim Geo(p), p = \frac{1}{75}$'
plt.title(title_str)
plt.axvline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.axhline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.vlines(the_x, 0, the_y, color='red', linestyles='-', lw=2)

Panel = Figure.add_subplot(2,2,2)
the_x = pmf_z['z'] 
the_y = pmf_z['Pr_z'] 
plt.xlim(x_lims)
plt.ylim(y_lims)
plt.xlabel('z (in Days)')
plt.ylabel('PMF $p(z)$')
title_str = r'PMF : $Z = X + Y$, $X,Y$ independent'
plt.title(title_str)
plt.axvline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.axhline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.vlines(the_x, 0, the_y, color='blue', linestyles='-', lw=2)

Panel = Figure.add_subplot(2,2,3)

LB_y = 0
UB_y = 1
off_set_y = 0.1*(UB_y-LB_y)
y_lims = (LB_y-off_set_y,UB_y+off_set_y)

the_x = x
the_y = cdf_x
plt.xlim(x_lims)
plt.ylim(y_lims)
plt.xlabel('x (in Days)')
plt.ylabel('CDF $F(x) = Pr(X \leq x)$')
title_str = r'CDF : $X$ or $Y \sim Geo(p), p = \frac{1}{75}$'
plt.title(title_str)
plt.axvline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.axhline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.plot(the_x,the_y, color='red', lw=1)
plt.axhline(1, lw=1, ls = '--',color = 'lightgray',alpha=0.5)

x_values = [101,101]
y_values = [0,the_y[100]]
plt.plot(x_values,y_values, color='black', lw=0.5, ls=':')
x_values = [0,101]
y_values = [the_y[100],the_y[100]]
plt.plot(x_values,y_values, color='black', lw=0.5, ls=':')

text_str = f'{the_y[100]:4.3f}'
plt.text(5,the_y[100],text_str,color = 'black', size = 8,horizontalalignment  = 'left', verticalalignment  = 'bottom')
    
Panel = Figure.add_subplot(2,2,4)

LB_y = 0
UB_y = 1
off_set_y = 0.1*(UB_y-LB_y)
y_lims = (LB_y-off_set_y,UB_y+off_set_y)

the_x = pmf_z['z'] 
the_y = pmf_z['Cum_Pr_z'] 
plt.xlim(x_lims)
plt.ylim(y_lims)
plt.xlabel('x (in Days)')
plt.ylabel('CDF $F(z) = Pr(Z \leq z)$')
title_str =r'CDF : $Z = X + Y$, $X,Y$ independent'
plt.title(title_str)
plt.axvline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.axhline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.plot(the_x,the_y, color='blue', lw=1)
plt.axhline(1, lw=1, ls = '--',color = 'lightgray',alpha=0.5)

x_values = [101,101]
y_values = [0,the_y[99]]
plt.plot(x_values,y_values, color='black', lw=0.5, ls=':')
x_values = [0,101]
y_values = [the_y[99],the_y[99]]
plt.plot(x_values,y_values, color='black', lw=0.5, ls=':')

text_str = f'{the_y[99]:4.3f}'
plt.text(5,the_y[99],text_str,color = 'black', size = 8,horizontalalignment  = 'left', verticalalignment  = 'bottom')

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