# -*- coding: utf-8 -*-
"""
Created on Sat Jan  8 15:19:57 2022

@author: Rene
"""

import numpy as np
import matplotlib.pyplot as plt

# Sampling from a Bernouli Distribution with probaility of succes equal to p
the_p = 0.75
m = 100
sample = np.random.binomial(size=m, n=1, p = 0.75)

# Transforming the Uniform Sample to a discrete distribution Sample
x_points = [0,1]
prob_points = [1-the_p,the_p]
y_points = np.zeros(2)
y_points[0] = len(sample[sample == x_points[0]])/m
y_points[1] = len(sample[sample == x_points[1]])/m

# Plotting the Uniform PDF and CDF

plt.rc('font', family='serif', size='12')
Figure = plt.figure(1)
Figure.set_figwidth(9)
Figure.set_figheight(5)
Figure.subplots_adjust(left=0.1, bottom=0.1, right=0.95, top=0.85, hspace=0.1, wspace=0.25)

#***************************************************************************************
LB_x = -0.25
UB_x = 1.125
off_set_x = 0.05*(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.1*(UB_y-LB_y)
y_lims = (LB_y-off_set_y,UB_y+off_set_y)

Panel = Figure.add_subplot(1,1,1)
plt.xlim(x_lims)
plt.ylim(y_lims)
plt.xlabel('x')
plt.ylabel('Probaility Mass Function $p(x)$')
text_str = r'PMF - $Ber(p), p = $'+ f'{the_p:0.2f}'
plt.title(text_str)
plt.axvline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.axhline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.axhline(prob_points[0], lw=1, ls = '--',color = 'darkgray',alpha=0.5)
plt.axhline(prob_points[1], lw=1, ls = '--',color = 'darkgray',alpha=0.5)

x_values = [0, 0]
y_values = [0, the_p]
#plt.plot(x_values, y_values,color='black',linestyle='-')
plt.scatter(x_points,prob_points,color = 'red',alpha=0.5,s=100)

text_str = r'sample size = '+ f'{m:0.0f}'
plt.text(0.025,0.9,text_str,horizontalalignment = 'left',verticalalignment = 'center', color = 'red')

for i in np.arange(0,2):
    if (i==0):
        text_str = r'$1 - p = $'+ f'{prob_points[i]:0.2f}'
    else:
        text_str = r'$p = $'+ f'{prob_points[i]:0.2f}'
    plt.text(x_points[i]+0.025,prob_points[i],text_str,horizontalalignment = 'left',verticalalignment = 'center', color = 'green')
    
    x_values = [x_points[i], x_points[i]]
    y_values = [0, prob_points[i]]
    plt.plot(x_values, y_values,color='red',lw=5,alpha=0.5)
    
    x_values = [x_points[i], x_points[i]]
    y_values = [0, y_points[i]]
    plt.plot(x_values, y_values,color='blue',lw=2,alpha=0.5)
    plt.scatter(x_points,y_points,color = 'blue',alpha = 0.5)
    
    text_str = r'$\widehat{p} = $'+ f'{y_points[i]:0.3f}'
    plt.text(x_points[i]-0.025,y_points[i],text_str,horizontalalignment = 'right',verticalalignment = 'center', color = 'blue')
    
plt.savefig('Bernouli_Sampler.png', dpi=300)