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


######################################################################################
# Los valores que siguen deben ser reemplazados por los del problema que corresponda
# Y las cantidades a graficar deben ser adaptadas al problema en cuestion
# acá se presenta como ejemplo el caso donde se quiere graficar el cuadrado de Y vs el cuadrado de X
# y luego ajustarlos por cuadrados minimos, antes y despues de elevar al cuadrado.
magnitud_en_X = np.array([10.5, 12.4, 14.7, 20.1, 23.1, 25.6, 30.0, 31.3, 33.5, 39.9]) #cm/s
magnitud_en_X = magnitud_en_X/100 # de cm/s a m/s

masa = 0.1 # kg
magnitud_en_x_al_cuadrado = magnitud_en_X**2*masa # kg m2/s2 = N m
#print(magnitud_en_x_al_cuadrado) 

# error_magnitud_en_X = np.array([0.5, 0.5, 0.5, 1.0, 1.0, 1.0, 1.5, 1.5, 1.5, 2.0])/100 #cm/s
# error_magnitud_en_x_al_cuadrado = 2*error_magnitud_en_X*magnitud_en_X
#print(error_magnitud_en_x_al_cuadrado)

magnitud_en_y = np.array([3.30, 4.10, 5.13, 6.92, 7.55, 8.43, 9.76, 10.32, 11.4, 13.28]) #N
magnitud_en_y_al_cuadrado = magnitud_en_y**2

error_magnitud_en_y = np.array([0.08, 0.10, 0.12, 0.22, 0.29, 0.35, 0.49, 0.53, 0.61, 0.86]) #N
error2_magnitud_en_y = 2*error_magnitud_en_y*magnitud_en_y
#print(error2_magnitud_en_y) 

######################################################################################

m, b, varm, varb = ajustelineal(magnitud_en_X,magnitud_en_y,error_magnitud_en_y)
print("Resultados para Y vs X")
print("Pendiente: {:.6f} +- {:.6f}".format(m, np.sqrt(varm)))
print("Ordenada : {:.6f} +- {:.6f}".format(b, np.sqrt(varb)))

chi2_experimental = chicuadrado(magnitud_en_X,magnitud_en_y,error_magnitud_en_y,m,b)
dof=len(magnitud_en_X)-2
print("Chi cuadrado/dof : {:.6f} / {:.0f}".format(chi2_experimental, dof))

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

m2, b2, varm2, varb2 = ajustelineal(magnitud_en_x_al_cuadrado,magnitud_en_y_al_cuadrado,error2_magnitud_en_y)
print("Resultados para Y cuadrado vs X cuadrado")
print("Pendiente: {:.6f} +- {:.6f}".format(m2, np.sqrt(varm2)))
print("Ordenada : {:.6f} +- {:.6f}".format(b2, np.sqrt(varb2)))

chi2_experimental2 = chicuadrado(magnitud_en_x_al_cuadrado,magnitud_en_y_al_cuadrado,error2_magnitud_en_y,m2,b2)
dof2=len(magnitud_en_x_al_cuadrado)-2
print("Chi cuadrado/dof : {:.6f} / {:.0f}".format(chi2_experimental2, dof2))

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

plt.figure(figsize=(8,5))
plt.errorbar(magnitud_en_x_al_cuadrado,magnitud_en_y_al_cuadrado,yerr=error2_magnitud_en_y,marker='h',ls='None', color='b', label='Datos')
plt.plot(magnitud_en_x_al_cuadrado,m2*magnitud_en_x_al_cuadrado+b2,color='red', label='Ajuste')
plt.tick_params(direction='in')
plt.xlabel("V$^{2}$ x Masa [N m]",fontsize=20)
plt.ylabel("F$^{2}$ [N$^{2}$]",fontsize=20)
plt.grid()
plt.show()