import math
import pandas as pd
import numpy as np
import matplotlib
import matplotlib.pyplot as plt

from astropy import stats
from mpl_toolkits.mplot3d import Axes3D
from scipy.optimize import curve_fit
def sqrt_N(x, c):
    return c/np.sqrt(x)

# Entrega el valor de rms para la observación de los datos entregados
def obtain_obs_error(data):
    data_noise = stats.sigma_clip(data, sigma=3)
    #plt.plot(np.linspace(-187.228, 137.228, data.size), data)
    #plt.plot(np.linspace(-187.228, 137.228, data.size), data_noise,'.')
    #plt.show()
    return data_noise.std()

# Grafica la disminución de rms al hacer mas observaciones, comparando el ruido de la observación central
# al ir sumando observaciones
def test_rms_decrease():
    min_size = 462
    number_of_obs = 12

    # Guarda todas las observaciones en un arreglo
    all_obs = []
    for i in range(1,4):
        for j in range(0,5):
            all_obs.append(get_data('grupo'+ str(i)+ '/sdf_6'+str(j)+ '_6'+ str(j)).Temperature.values[0:min_size])

    # Calcula el rms la suma de las observacioes y las guarda en un arreglo
    obs_rms = []
    current_sum = np.zeros(min_size)
    for i in range(0,number_of_obs):
        current_sum = current_sum + all_obs[i]
        obs_rms.append(obtain_obs_error(current_sum))


    n_array = np.linspace(1,number_of_obs,number_of_obs)

    p0 = [0.01]
    coeff, var_matrix = curve_fit(sqrt_N, n_array, obs_rms, p0=p0)
    obs_rms_fit = sqrt_N(np.linspace(0.9,number_of_obs), coeff[0])

    fig = plt.figure()
    plt.plot(n_array, obs_rms, '.')
    plt.plot(np.linspace(0,number_of_obs), obs_rms_fit)
    plt.title(r'Dismiución de $\Delta T_{out}$ al aumentar el número de observaciones')
    plt.xlabel('Número de observaciones')
    plt.ylabel(r'$\Delta T_{out}$')

    plt.show()
test_rms_decrease()