# -*- coding: utf-8 -*-
"""
Created on Thu Feb  3 18:47:48 2022

@author: Rene
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import gamma
from scipy.optimize import brentq
from math import trunc

def Pr_Credibility_Interval(n,the_lambda,HalfWidth):
    mu = 1/the_lambda
    LB = mu-HalfWidth
    UB= mu+HalfWidth
    value=gamma.cdf(UB,a = n,scale = 1/(n*the_lambda))-gamma.cdf(LB,a = n,scale = 1/(n*the_lambda))
    return(value)

def solve_n(p,n,the_lambda,HalfWidth,ub_n):
    
    def f(n):
        value = p-Pr_Credibility_Interval(n,the_lambda,HalfWidth)
        return value
    
    the_sol = brentq(f, a=0.0001, b=ub_n) 
    return the_sol



the_lambda = 0.5
mu = 1/the_lambda
half_width = 0.5
p = 0.90
start_n = 1
UB_n = 20001
n_sol = solve_n(p,start_n,the_lambda,half_width,UB_n)

n_values = np.arange(0.001,max(trunc(1.1*n_sol),125))
Pr_Cred_Int = Pr_Credibility_Interval(n_values,the_lambda,half_width)


# Plotting a four panel figure

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.5, wspace=0.25)

# Adding a panel to the figure
Panel = Figure.add_subplot(1,1,1)  
LB_x = 0
UB_x = max(trunc(1.1*n_sol),125)
off_set_x = 0.025*(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 = n_values
the_y = Pr_Cred_Int
plt.xlim(x_lims)
plt.ylim(y_lims)
plt.xlabel('x')
title_str = r'$\Pr[\bar{X}_n \in (\mu-$'+ f'{half_width:2.1f}'+'$,\mu+$+'f'{half_width:2.1f}'+'$)]$'
plt.ylabel(title_str, size = 10)
plt.title(title_str, size = 12)
plt.axvline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.axhline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.axhline(1, lw=2, ls = '--',color = 'lightgray',alpha=0.5)
plt.plot(the_x,the_y, color='red', lw=1)
x_values = [n_sol,n_sol]
y_values = [0,Pr_Credibility_Interval(n_sol,the_lambda,half_width)]
plt.plot(x_values,y_values, color='black', lw=1, ls=':')
x_values = [0,n_sol]
y_values = [Pr_Credibility_Interval(n_sol,the_lambda,half_width),Pr_Credibility_Interval(n_sol,the_lambda,half_width)]
plt.plot(x_values,y_values, color='black', lw=1, ls=':')


text_str = f'{p:2.2f}'
plt.text(1,0.9,text_str,color = 'red', size = 10,horizontalalignment  = 'left', verticalalignment  = 'bottom')
text_str = f'{n_sol:2.2f}'
plt.text(n_sol-1,0.025,text_str,color = 'red', size = 10,horizontalalignment  = 'right', verticalalignment  = 'bottom')


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