# -*- 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
import Dist_Library as dl

def pmf_unif(x,min,max):
  value= np.zeros(len(x))
  for i in np.arange(0,len(x)):
      if ((x[i]>=min) & (x[i]<=max)):
          value[i] = 1/(max-min+1)
  return(value)

def pmf_convolution_unif(min,max):
    s = np.arange(2*min,2*max+1)
    Pr_s = np.zeros(len(s))
    y = np.arange(min,max+1)
    Pr_y = pmf_unif(y,min,max)
    for i in np.arange(0,len(s)):
        x = s[i]-y
        Pr_x = pmf_unif(x,min,max)
        Pr_s[i] = sum(Pr_x*Pr_y)
    value = pd.DataFrame({'s':s, 'Pr_s':Pr_s})
    return(value)


min = 1
max = 6
s = np.arange(2*min,2*max)
y = np.arange(min,max+1)
pmf_s = pmf_convolution_unif(min,max)

# 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.1, wspace=0.3)
off_set_x = 10

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

the_x = pmf_s['s'] 
the_y = pmf_s['Pr_s'] 
plt.xlim(x_lims)
plt.ylim(y_lims)
plt.xlabel('s')
plt.ylabel('PMF $p(s)$')
title_str = r'PMF : $S = 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='red', linestyles='-', lw=2)
plt.scatter(the_x, the_y, color='red', lw=2, s = 5)

Panel = Figure.add_subplot(1,2,2)
LB_x = 1
UB_x = 13
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 = 1
off_set_y = 0.05*(UB_y-LB_y)
y_lims = (LB_y-off_set_y,UB_y+off_set_y)
the_x = pmf_s['s'] 
the_y = pmf_s['Pr_s'] 
plt.xlim(x_lims)
plt.ylim(y_lims)
plt.xlabel('s')
plt.ylabel('CDF $F(s) = Pr(S \leq s)$')
title_str = r'PMF : $S = 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)
dl.plot_discrete_cdf(Panel,the_x,the_y,'red',10)

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