import numpy as np
from numpy import random
from scipy.optimize import curve_fit
import matplotlib.pyplot as plt
from scipy.stats import chi2

def lineal(x,m,b):
	return m*x+b

def ajustelineal(x,y,dy):
	w=1/dy**2
	SX=np.sum(x*w)
	SY=np.sum(w*y)
	SX2=np.sum(w*(x**2))
	SXY=np.sum(w*x*y)
	SW=np.sum(w)
	Xmed=SX/SW
	Ymed=SY/SW
	Varx=SX2/SW-Xmed**2
	delta=(SW**2)*Varx
	m=(SXY*SW-SX*SY)/delta
	b=(SX2*SY-SX*SXY)/delta
	Varm=np.sum(w*((y-m*x-b)**2))/((len(x)-2)*SW*Varx )
	Varb=Varm * SX2 / SW
	return m,b,Varm,Varb  # Notar que calcula las varianzas y no los sigmas

def chicuadrado(x,y,dy,m,b):
	w=1/dy**2
	chicuadrado=np.sum(w*(y-m*x-b)**2)
	return chicuadrado

def calcular_p_valor(realizacion_chi2, dof):
    p_valor = 1 - chi2.cdf(realizacion_chi2, dof)
    return p_valor


n=10							#cantidad de longitudes medidas
longitudes=np.linspace(1.5,2.5,n)	#[cm]

g=9.81
R=4*3.14159**2/g	
T2=R*longitudes					#esperanza de la variable aleatoria Periodo al cuadrado

dT2=np.zeros(n)+0.2				#incerteza en la medición del periodo al cuadrado

T2med=np.zeros(n)				#Valores medidos (Montecarleados) de los periodos al cuadrado
T2med=random.normal(T2,dT2)

m, b, varm, varb = ajustelineal(longitudes,T2med,dT2)
print("Pendiente: {:.2f} +- {:.2f}".format(m, np.sqrt(varm)))
print("Ordenada : {:.2f} +- {:.2f}".format(b, np.sqrt(varb)))

chi2_experimental = chicuadrado(longitudes,T2med,dT2,m,b)
dof=n-2
print("Chi cuadrado/dof : {:.2f} / {:.0f}".format(chi2_experimental, dof))

p_valor = calcular_p_valor(chi2_experimental, dof)
print("El valor p asociado es: {:.4f}".format(p_valor))

plt.figure(figsize=(8,5))
plt.errorbar(longitudes,T2med,yerr=dT2,marker='h',ls='None', color='b', label='Datos')
plt.plot(longitudes,m*longitudes+b,color='red', label='Ajuste')
plt.tick_params(direction='in')
plt.xlabel("Longitudes [cm]",fontsize=20)
plt.ylabel("$Periodo^{2}$ [s$^{2}$]",fontsize=20)
plt.show()