# -*- coding: utf-8 -*-
"""
Created on Thu Jul 21 15:28:41 2022

@author: Rene
"""

import numpy as np
from scipy.stats import norm
import matplotlib.pyplot as plt

def solve_quadratic(a,b,c):
    D = b*b-4*a*c
    x = np.zeros(2)
    if (D>=0):
        x[0] = (-b-D**0.5)/(2*a)
        x[1] = (-b+D**0.5)/(2*a)
    return(x)

def parabola(x,a,b,c):
  value = a*x*x+b*x+c
  return(value)

n = 30
p = 0.5
voter_sample = np.random.binomial(1, p, size=n)
x = sum(voter_sample)
conf_level = 0.90
alpha = 1-conf_level
z_alpha = norm.ppf(1-alpha/2,0,1)

# Central Limit Theorem method Confidence Interval for proportion

p_hat = x/n
sd_p_hat = (p_hat*(1-p_hat)/n)**0.5
half_width = z_alpha*sd_p_hat
conf_LB = p_hat-half_width
conf_UB = p_hat+half_width
mu_in_conf = ((conf_LB <= p) & (conf_UB  >= p))

# Wilson method Confidence Interval for proportion

a = (z_alpha**2)/n+1
b = -(2*x+(z_alpha**2))/n
c = (x/n)**2
p_sol = solve_quadratic(a,b,c)
p_in_conf = ((p_sol[0] <= p) & (p_sol[1] >= p))
if (p_in_conf==1):
    color = "green3"
else:
    color = "red"
    
# Evaluating parabola points
LB = -0.05
UB = 1.05
x_points = np.linspace(LB,UB,100)
y_points = parabola(x_points,a,b,c) 

# Plotting the analysis results    
off_set_x = 0.05
off_set_y = 0.0025
max_y = 0.4
    
# Setting the attributes for the Figure     
plt.rc('font', family='serif', size='10')
plt.figure(figsize=(10, 5))
Figure = plt.figure()
Figure.subplots_adjust(hspace=0.4, wspace=0.4)
x_lims = (LB-off_set_x,UB+off_set_x)
y_lims = (min(min(y_points),-0.0175)-off_set_y,max(-min(y_points),0.0175)+off_set_y)

min(min(y_points),-0.0175)

plt.axvline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.axhline(0, lw=2, ls = '-',color = 'lightgray',alpha=0.5)
plt.xlim(x_lims)
plt.ylim(y_lims)
plt.xlabel('p')
plt.ylabel('y')
plt.plot(x_points,y_points,ls='-',lw=0.5,label='parabola',color='blue')

plt.axvline(1, lw=1, ls = '--',color = 'lightgray',alpha=0.5)   
plt.axvline(p, lw=1, ls = ':',color = 'red',alpha=0.5)    
   
x_values = [p_sol[0], p_sol[0]]
y_values = [0, -0.01]
plt.plot(x_values,y_values,lw=0.5,ls=':',color='black') 

x_values = [p_sol[1], p_sol[1]]
y_values = [0, -0.01]
plt.plot(x_values,y_values,lw=0.5,ls=':',color='black')   

x_values = [x/n, x/n]
y_values = [0, -1]
plt.plot(x_values,y_values,lw=0.5,ls='--',color='black')  
text_str = '$x/n = $'+ f'{x/n:2.2f}'
plt.text(x/n,0,text_str,horizontalalignment = 'center', verticalalignment = 'bottom')

text_str = 'n ='+ f'{n:2.0f}'
plt.text(0,-min(y_points),text_str,horizontalalignment = 'left', verticalalignment = 'top', color = 'red')

text_str = 'x ='+ f'{x:2.0f}'
plt.text(0,-min(y_points)-0.004,text_str,horizontalalignment = 'left', verticalalignment = 'top', color = 'red')

# Plotting Wilson Confidence Interval
if p_in_conf:
    conf_color = 'green'
else:
    conf_color = 'red'
half_width = (p_sol[1] - p_sol[0])/2
mid_point = (p_sol[1] + p_sol[0])/2
plt.errorbar(mid_point, 0, xerr = half_width, ecolor=conf_color, 
             elinewidth=1, capsize=3,fmt='',ms=3,mfc=conf_color,mec=conf_color, label = 'Wilson Method', c = conf_color);
plt.scatter(x/n,0,color=conf_color,s = 10)
text_str = f'{p_sol[0]:2.2f}'
plt.text(p_sol[0],-0.01,text_str,horizontalalignment = 'center', verticalalignment = 'top', color = conf_color)
text_str = f'{p_sol[1]:2.2f}'
plt.text(p_sol[1],-0.01,text_str,horizontalalignment = 'center', verticalalignment = 'top', color = conf_color)

# Plotting CLT Confidence Interval
if mu_in_conf:
    conf_color = 'green'
else:
    conf_color = 'red'
half_width = (conf_UB - conf_LB)/2
eb1 = plt.errorbar(x/n, -0.015, xerr = half_width, ecolor=conf_color,
             elinewidth=1, capsize=3,fmt='-o',ms=3,mfc=conf_color,mec=conf_color, label = 'CLT Method', c = conf_color);
eb1[-1][0].set_linestyle('--')
 
text_str = f'{conf_LB:2.2f}'
plt.text(conf_LB-0.01,-0.015,text_str,horizontalalignment = 'right', verticalalignment = 'center', color = conf_color)
text_str = f'{conf_UB:2.2f}'
plt.text(conf_UB+0.01,-0.015,text_str,horizontalalignment = 'left', verticalalignment = 'center', color = conf_color)

plt.legend(loc='lower right')

alpha_str = f'{100*conf_level:2.0f}' + "% Confidence Interval proportion p"
Figure.suptitle(alpha_str,size='14')
plt.savefig("Wilson_Confidence_Interval_Proportion.png", dpi=300)