###############################################################################
# CARGAMOS LAS LIBRERIAS DE FUNCIONES HECHAS QUE VAMOS A USAR
###############################################################################
import numpy as np
import xlrd #PARA CARGAR ARCHIVOS DE EXCEL
import matplotlib.pyplot as plt
import matplotlib.mlab as mlab
# Esto es necesario para realizar los ajustes lineales
from scipy.optimize import curve_fit
#https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.curve_fit.html

###############################################################################
# DEFINIMOS NOSOTROS OTRAS DE LAS FUNCIONES QUE VAYAMOS A USAR
###############################################################################
def abrir_columna_excel(nombre_archivo,columna):
  wb = xlrd.open_workbook(nombre_archivo) 
  Hoja1 = wb.sheet_by_index(0) 
  N_datos=len(Hoja1.col_values(columna))
  datos=[]
  for i in range(1,N_datos):
    if type(Hoja1.cell_value(i, columna))==float:
      datos.append(Hoja1.cell_value(i, columna))
  return datos

# Esta función es necesaria para calcular el R^2
def R_cuadrado(y,y_ajuste):
    ss_res = np.sum((y - y_ajuste) ** 2) # Suma de los residuos al cuadrado
    ss_tot = np.sum((y - np.mean(y)) ** 2) # Varianza
    R2=1-(ss_res/ss_tot) # R cuadrado
    return R2
###############################################################################
# ESCRIBIMOS LA SECUENCIA DE PROCESOS DE NUESTRO PROGRAMA
###############################################################################
# CARGAMOS LAS COLUMNAS DE PERIODOS Y ERRORES
nombre_archivo = ("C:\\Users\\Usuario\\Dropbox\\Clase Estadística\\periodos.xlsx")
columna=0
largos=abrir_columna_excel(nombre_archivo,columna)
columna=1
periodos=abrir_columna_excel(nombre_archivo,columna)
columna=2
errores=abrir_columna_excel(nombre_archivo,columna)

# Definimos x, y y el error en y para ajustar, como arrays
x = np.asarray(largos)
y = np.asarray(periodos)**2
yerr = np.asarray(errores)*2*np.asarray(periodos)

# Definimos la funcion para ajustar ('x, parametros: formula')
funcion_lineal = lambda x, ordenada, pendiente: pendiente * x + ordenada

# ajusto la funcion lineal, con los datos x,y y el error en y.
parametros, pcov = curve_fit(funcion_lineal, x, y, sigma = yerr,absolute_sigma=True)

# obtengo los resultados del ajuste
ordenada=parametros[0]
pendiente=parametros[1]

# las incertezas de los parametros son la raiz de la diagonal de la matriz de covarianza
sigmas = np.sqrt([pcov[0,0],pcov[1,1]])
error_ordenada=sigmas[0]
error_pendiente=sigmas[1]

print('Ordenada al origen: '+str(ordenada)+' ± '+str(error_ordenada))
print('Pendiente: '+str(pendiente)+' ± '+str(error_pendiente))

#Calculo el R cuadrado
y_ajuste=pendiente*x+ordenada
R2=R_cuadrado(y,y_ajuste)

# graficamos
plt.errorbar(x,y,yerr=yerr,color='k',fmt='o')
plt.plot(x,y,'ob',label='Datos')#,10000,fmt='-o',color='g')
plt.xlabel('Largos')
plt.ylabel('Periodos')
plt.plot(x,pendiente*x+ordenada,'r',label='Ajuste')
plt.title('$R^2$ = '+str(R2))
plt.legend()
plt.show()