import numpy as np
import matplotlib.pyplot as plt
import pyfits
from scipy.stats import sigmaclip as sic
import scipy.signal as sp

'''Antes de partir, para correr el codigo, se asigna a,b para el masterflat y
espectro respectivamente, esto se ingresa al metodo yay que luego devuelve la
longitud de onda y flujo respectivamente. Estos dos valores son luego
ingresados en el metodo Mask, junto con el nombre del archivo con la mascara y
la velocidad que se quiere utilizar. Recuerdo aca tambien que la seccion final
de Mask no logro funcionar correctamente, por lo que se obtienen valores
demasiado pequenos para las velocidades buscadas'''

#Primero, creo un metodo para leer los datos y corregir con el masterflat:

def corr (mf, spec):
    sp = pyfits.getdata(spec)
    ps = sp
    fl = pyfits.getdata(mf)

    #Corrijo el flujo del espectro con el del flat dividiendo:

    for i in range(1024):
        for j in range(41):
            ps[1][j][i] = (sp[1][j][i])/(fl[j][1][i])

    #Veo en donde se hicieron divisiones por 0 (quedan nan)
            #y elimino esos valores

    nanas = np.isnan(ps)
    wav = []
    flux = []
    for i in range(41):
        wav.append(ps[0][i][-nanas[1][i]])
    for i in range(41):
        flux.append(ps[1][i][-nanas[1][i]])
    wav = np.array(wav)
    flux = np.array(flux)

    #En el final de los datos queda todo muy raro. Se cortaran los datos
    # al final para poder realizar bien el ajuste polinomial y para
    #coincidir con la mascara. Ademas, todo seve horrible al principio,
    #por lo que se cortara tambien.
    Wav = []
    Flux = []
    for i in range(41):
        if wav[i][0] > 4500 and wav[i][1022]<6800:
            Wav.append(wav[i])
            Flux.append(flux[i])
    Flux = np.array(Flux)
    Wav = np.array(Wav)


    return Wav,Flux

#A continuacion, procedemos a realizar el fiteo. Para esto, se usa un polinomio
# de grado 3

def fiteo(wav,flux):
    coef = []
    cosa = []
    for i in range(len(flux)):
        poli = np.polyfit(wav[i],flux[i], 3)

        y = np.poly1d(poli)
        x = wav[i]
        f = y(x)
        cosa.append(f)

    return np.array(cosa)

#Ahora, hago elimino rayos cosmicos usando sigma clipping:

def sick(Wav,Flux):

    sig = []
    FL = []
    WV = []
    end1 = []
    end2 = []
    for i in range (len(Flux)):
        mad = np.median(np.abs(Flux[i]-np.median(Flux[i])))
        sig.append(1.4826*mad)
    for i in range (len(Flux)):
        aux1=[]
        aux2=[]
        for j in range(1023):
            if (np.abs(Flux[i][j]-np.median(Flux[i]))<3*sig[i]):
                aux1.append(Flux[i][j])
                aux2.append(Wav[i][j])

        if aux1:
            FL.append(aux1)
            WV.append(aux2)
    return np.array(WV),np.array(FL)

#Finalmente, unimos todo en un metodo que plotea lo que queremos. Es decir,
# la curva ajustada para cada orden y luego con la aplicacion de sigma clipping

def yay (mf,spec):
    wav, flux = corr(mf,spec)
    f = fiteo(wav,flux)
    flux  = flux/f
    c, d = sick(wav,flux)
    for i in range(len(c)):
        plt.plot(c[i], d[i], linewidth='0.1')
    return c,d


#Ahora, creo el metodo de aplicacion de la mascara (antes de moverla)
#Es decir, es para un valor dado de cada cosa
def apmask(wavu,fluxu,masc,wind):
    f = 0
    fl = 0
    for i in range(10):
        if wavu<masc[i][0]:
            if wavu+wind>masc[i][0]:
                fl = masc[i][2]*fluxu*(wavu+wind-masc[i][0])/(2*wind)
        if wavu>masc[i][0]:
            if wavu-wind<masc[i][0]:


                fl = masc[i][2]*fluxu*(wavu+wind-masc[i][0])/(2*wind)

            if wavu+wind<masc[i][1] and wavu-wind>masc[i][0]:
                fl = masc[i][2]*fluxu

            if wavu+wind>masc[i][1] and wavu-wind<masc[i][1]:
                fl = masc[i][2]*fluxu*(masc[i][1]-(wavu-wind))/(2*wind)
        if wavu>masc[i][1]:
            if wavu-wind<masc[i][1]:
                masc[i][2]*fluxu*(masc[i][1]-(wavu-wind))/(2*wind)
        f = f + fl
    return f

#Luego, creamos un metodo que mueva la mascara (aca, vc = v/c y le otorgamos
#un valor arbitrario, por ejemplo, 0.1:

def movmask(masc,vc):
    masc = np.array(masc)
    aux = masc.T
    delam = aux[0]*vc
    aux[0] = aux[0]+delam
    aux[1] = aux[1]+delam
    masc = aux.T
    return masc

#Ahora, queremos aplicar la mascara:
def Mask (wav,flux,masc,vc):
    mask = np.loadtxt(masc)
    aux = []
    for i in range(len(mask)):
        if mask[i][0] > 4500:
            aux.append(mask[i])
    mask = aux



#Necesito crear una ventana para las longitudes de onda de cada pixel:
    aux1 = 100
    wind = []
    for i in range(len(wav)):
        for j in range(len(wav[i])-1):
            win = np.abs(wav[i][j+1]-wav[i][j])
            if win < aux1:
                aux1 = win

        wind.append(aux1)
#Con esto, tenemos una ventana minima aproximada para cada pixel en los ordenes
#Esto lo usaremos para ir aplicando la mascara a cada flujo en cada orden, para
#luego sumar todo:
    suma = []
    x=0
    while mask[0][0] < 6800:
        cosa = 0

        for i in range(len(wav)):
            for j in range(len(wav[i])):
                cosa = cosa + apmask(wav[i][j],flux[i][j], mask, wind[i])
        suma.append(cosa)
        mask = movmask(mask,vc)
        #print('Vuelta! Valor: ' + str(cosa))
        x = x+1
        if x%3 == 0:
            print('Valor: ' + str(cosa))
            x = 0
    return suma