# -*- coding: utf-8 -*-
"""
Created on Wed Mar  4 18:10:08 2020

@author: Johan Rene van Dorp
"""
import numpy as np
import matplotlib.pyplot as plt
import Dist_Library as dl
from scipy.stats import binom
from scipy.stats import geom


prob = 1/2
x_coin = np.arange(2)
p_coin = np.arange(1)
p_coin =[1-prob,prob]

prob = 1/6
x_die = np.arange(1,7)
p_die = np.repeat(prob,6)

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

# Adding the left panel with the theoretical pmf
Left_Panel = PMF_figure.add_subplot(1,2,1)
plt.scatter(x_coin,p_coin,color='red')
plt.xlim((-0.5,1.5))
plt.ylim((-0.05,1))
plt.xlabel('x')
plt.ylabel('Probability Mass Function $p(x)$')
plt.title('PMF - Coin Toss Example')
plt.axvline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.axhline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.vlines(x_coin, 0, p_coin, color='red', linestyles='-', lw=2)

# Adding the right panel with the empirical pmf
Right_Panel = PMF_figure.add_subplot(1,2,2)
plt.scatter(x_die,p_die,color='steelblue')
plt.xlim((-0.5,7))
plt.ylim((-0.05,1))
plt.xlabel('x')
plt.ylabel('Probability Mass Function $p(x)$')
plt.title('PMF - Die Example')
plt.axvline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.axhline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.vlines(x_die, 0, p_die, color='steelblue', linestyles='-', lw=2)

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

# Creating the cumulative distribution figure
plt.rc('font', family='serif', size='10')
CDF_figure = plt.figure()
CDF_figure.set_figwidth(7)
CDF_figure.set_figheight(3.5)
CDF_figure.subplots_adjust(left=0.1, bottom=0.15, right=0.95, top=0.9, hspace=0.1, wspace=0.4)

# Adding the left panel with the theoretical pmf
Left_Panel = CDF_figure.add_subplot(1,2,1)
dl.plot_discrete_cdf(Left_Panel,x_coin,p_coin,'red',1)
plt.xlim((-0.5,1.5))
plt.ylim((-0.05,1.1))
plt.xlabel('x')
plt.ylabel('Cumulative Distr. Function $F(x) = Pr(X \leq x)$')
plt.title('CDF - Coin Toss Example')
plt.axvline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.axhline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)

# Adding the right panel with the empirical pmf
Right_Panel = CDF_figure.add_subplot(1,2,2)
dl.plot_discrete_cdf(Right_Panel,x_die,p_die,'steelblue',2)
plt.xlim((-0.5,7))
plt.ylim((-0.05,1.1))
plt.xlabel('x')
plt.ylabel('Cumulative Distr. Function $F(x) = Pr(X \leq x)$')
plt.title('CDF - Die Example')
plt.axvline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.axhline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)

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

# Creating the binomial distribution figure
the_n = 10
the_prob = 0.25
x_binom = np.arange(0,the_n+1)
p_binom = binom.pmf(np.arange(0,the_n+1), the_n, the_prob)

plt.rc('font', family='serif', size='10')
Binom_figure = plt.figure()
Binom_figure.set_figwidth(7)
Binom_figure.set_figheight(3.5)
Binom_figure.subplots_adjust(left=0.1, bottom=0.15, right=0.95, top=0.9, hspace=0.1, wspace=0.4)

# Adding the left panel with the theoretical pmf
Left_Panel = Binom_figure.add_subplot(1,2,1)
plt.scatter(x_binom,p_binom,color='red')
plt.xlim((-0.5,10.5))
plt.ylim((-0.015,0.3))
plt.xlabel('x')
plt.ylabel('Probability Mass Function $p(x)$')
plt.title('PDF : $X \sim Binom(10,0.25)$')
plt.axvline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.axhline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.vlines(x_binom, 0, p_binom, color='red', linestyles='-', lw=2)

# Adding the right panel with the empirical pmf
Right_Panel = Binom_figure.add_subplot(1,2,2)
dl.plot_discrete_cdf(Right_Panel,x_binom,p_binom,'red',2)
plt.xlim((-0.5,10.5))
plt.ylim((-0.05,1.1))
plt.xlabel('x')
plt.ylabel('Cumulative Distr. Function $F(x) = Pr(X \leq x)$')
plt.title('CDF : $X \sim Binom(10,0.25)$')
plt.axvline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.axhline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)

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

# Creating the geometric distribution figure
the_n = 15
the_prob = 0.25
x_binom = np.arange(1,the_n+1)
p_binom = geom.pmf(x_binom, the_prob)

plt.rc('font', family='serif', size='10')
Geom_figure = plt.figure()
Geom_figure.set_figwidth(7)
Geom_figure.set_figheight(3.5)
Geom_figure.subplots_adjust(left=0.1, bottom=0.15, right=0.95, top=0.9, hspace=0.1, wspace=0.4)

# Adding the left panel with the theoretical pmf
Left_Panel = Geom_figure.add_subplot(1,2,1)
plt.scatter(x_binom,p_binom,color='steelblue')
plt.xlim((-0.5,15.5))
plt.ylim((-0.015,0.3))
plt.xlabel('x')
plt.ylabel('Probability Mass Function $p(x)$')
plt.title('PDF : $X \sim Geom(0.25)$')
plt.axvline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.axhline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.vlines(x_binom, 0, p_binom, color='steelblue', linestyles='-', lw=2)

# Adding the right panel with the empirical pmf
Right_Panel = Geom_figure.add_subplot(1,2,2)
dl.plot_discrete_cdf(Right_Panel,x_binom,p_binom,'steelblue',2)
plt.xlim((-0.5,15.5))
plt.ylim((-0.05,1.1))
plt.xlabel('x')
plt.ylabel('Cumulative Distr. Function $F(x) = Pr(X \leq x)$')
plt.title('CDF : $X \sim Geom(0.25)$')
plt.axvline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.axhline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)

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