import matplotlib
matplotlib.use("Agg") 
from astropy.io import fits as pyfits
import numpy as np
from numpy import median,sqrt,array,exp
import scipy
from scipy import signal,special,optimize,interpolate
import scipy.special as sp
import copy
import glob
import os
import matplotlib.pyplot as plt
import sys
from ceres3.utils import globalutils as GLOBALutils


def FileClassify(diri, log,mode='so',amps='a'):
        """
        Classifies all files in a directory and writes a night log of science images
        """
        if mode == 'so':
                md = 'star only'
        elif mode == 'ss':
                md = 'star+sky'
        else:
                md = 'Polarimetry'

        if amps == 'a':
                samps = 'a'
        elif amps == 'ab':
                samps = 'a,b'

        # define output lists
        simThAr_sci    = []
        simSky_sci     = []
        biases         = []
        bias_dates     = []
        flat_dates     = []
        flats          = []
        ThAr_ref       = []
        ThAr_ref_dates = []
        mode           = []

        f = open(log,'w')
        bad_files = []
        if os.access(diri+'bad_files.txt',os.F_OK):
                bf = open(diri+'bad_files.txt')
                linesbf = bf.readlines()
                for line in linesbf:
                        bad_files.append(diri+line[:-1])
                bf.close()

        all_files = glob.glob(diri+"*fits")

        jj = 0

        for archivo in all_files:
                jj+=1
                dump = False
                for bf in bad_files:
                        if archivo == bf:
                                dump = True
                                break

                if not dump:
                        h = pyfits.open(archivo)
                        #print archivo, h[0].header['EXPTYPE'], h[0].header['INSTMODE']
                        if md in h[0].header['INSTMODE'] and samps == h[0].header['AMPLIST']:
                                print(archivo, h[0].header['EXPTYPE'], h[0].header['INSTMODE'], h[0].data.shape, h[0].header['AMPLIST'])
                                if h[0].header['EXPTYPE'] == 'OBJECT':
                                        simThAr_sci.append(archivo)
                                        obname = h[0].header['OBJNAME']
                                        ra     = h[0].header['OBJRA']
                                        delta  = h[0].header['OBJDEC']
                                        airmass= h[0].header['AIRMASS']
                                        texp   = h[0].header['EXPTIME']
                                        date   = h[0].header['DATE-OBS'] + h[0].header['UTC-OBS']
                                        line = "%-15s %10s %10s %8.2f %4.2f %s %8s %s\n" % (obname, ra, delta, texp, airmass, h[0].header['EXPTYPE'], date, archivo)
                                        f.write(line)

                                elif h[0].header['EXPTYPE'] == 'BIAS':
                                        biases.append(archivo)
                                        mjd,mjd0 = mjd_fromheader(h)
                                        bias_dates.append(mjd)

                                elif h[0].header['EXPTYPE'] == 'FLAT':
                                        flats.append(archivo)
                                        mjd,mjd0 = mjd_fromheader(h)
                                        flat_dates.append(mjd)

                                elif h[0].header['EXPTYPE'] == 'COMPARISON':
                                        mjd, mjd0 = mjd_fromheader(h)
                                        ThAr_ref.append(archivo) 
                                        ThAr_ref_dates.append( mjd )

        f.close()
        biases, bias_dates = np.array(biases), np.array(bias_dates)
        flats, flat_dates  = np.array(flats), np.array(flat_dates)
        IS = np.argsort(bias_dates)
        biases, bias_dates = biases[IS], bias_dates[IS]
        IS = np.argsort(flat_dates)
        flats, flat_dates = flats[IS], flat_dates[IS]

        return biases, flats, ThAr_ref, simThAr_sci, ThAr_ref_dates

def mjd_fromheader(h):
    """
    return modified Julian date from header
    """
    
    datetu = h[0].header['DATE-OBS']
    timetu = h[0].header['UTC-OBS'].split(':')

    mjd0,mjd,i = GLOBALutils.iau_cal2jd(int(datetu[:4]),int(datetu[5:7]),int(datetu[8:10]))

    ho = int(timetu[0])
    mi = int(timetu[1])
    se = float(timetu[2])
    ut = float(ho) + float(mi)/60.0 + float(se)/3600.0
    mjd_start = mjd + ut/24.0

    secinday = 24*3600.0
    fraction = 0.5
    texp     = h[0].header['EXPTIME'] #sec

    mjd = mjd_start + (fraction * texp) / secinday

    return mjd, mjd0

def MedianCombine(ImgList, zero_bo=False, zero='MasterBias.fits'):
        """
        Median combine a list of images
        """
        if zero_bo:
                BIAS = pyfits.getdata(zero)

        n = len(ImgList)
        if n==0:
                raise ValueError("empty list provided!")

        h = pyfits.open(ImgList[0])[0]
        d = h.data
        d = OverscanTrim(d)
        if zero_bo:
                d -= BIAS
                d = np.round(d).astype('int')
        factor = 1.25
        if (n < 3):
                factor = 1
        ronoise = factor * h.header['RDNOISEA'] / np.sqrt(n)
        gain    = h.header['GAINA']
        if (n == 1):
                return d, ronoise, gain
        else:
                for i in range(n-1):
                        h = pyfits.open(ImgList[i+1])[0]
                        ot = OverscanTrim(h.data)
                        if zero_bo:
                                d = np.dstack((d,np.round((ot - BIAS)).astype('int')))
                        else:
                                d = np.dstack((d,np.round(ot).astype('int')))
        return np.median(d,axis=2), ronoise, gain

def OverscanTrim(d):
    """
    Overscan correct and Trim a refurbished FEROS image
    """
    os = d[:,2060:]
    overscan = np.median(os,axis=1)[:4608]
    newdata = np.zeros(d[:4608,:2049].shape)
    for i in range(len(overscan)):
        newdata[i,:] = d[i,:2049] - overscan[i]
    return newdata

def get_them(sc,exap,ncoef,mode=1,shift=0.):
        def gauss(params,x):
                amp1 = params[0]
                med1 = params[1]
                sig1 = params[2]
                g1 = amp1 * np.exp(-0.5*((x-med1)/sig1)**2)
                return g1
        def res_gauss(params,g,x):
                return g-gauss(params,x)

        def gauss2(params,x):
                amp1 = params[0]
                amp2 = params[1]
                med1 = params[2]
                med2 = params[3]
                sig1 = params[4]
                sig2 = params[5]
                g1 = amp1 * np.exp(-0.5*((x-med1)/sig1)**2)
                g2 = amp2 * np.exp(-0.5*((x-med2)/sig2)**2)
                return g1 + g2

        def res_gauss2(params,g,x):
                return g-gauss2(params,x)

        def fit_drift(ref,dif,oref):
                coef  = np.polyfit(ref,dif,2)

                residuals = dif - np.polyval(coef,ref)
                rms = np.sqrt(np.var(residuals))
                I = np.where(np.absolute(residuals)>3*rms)[0]
                cond = True
                if len(I)==0:
                        cond = False
                while cond:
                        im  = np.argmax(np.absolute(residuals))
                        dif = np.delete(dif,im)
                        ref = np.delete(ref,im)
                        coef = np.polyfit(ref,dif,2)
                        residuals = dif - np.polyval(coef,ref)
                        rms = np.sqrt(np.var(residuals))
                        I = np.where(np.absolute(residuals)>3*rms)[0]
                        if len(I)==0:
                                cond = False

                cdif = np.polyval(coef,oref)
                ref = oref + cdif
                return ref
        
        sc_or = sc.copy()  
        medc = int(.5*sc.shape[1])
        d = np.median(sc[:,medc-exap:medc+exap+1],axis=1)
        #plot(d)
        ejx = np.arange(len(d))

        ref = []
        cents = np.loadtxt('ref_ords.txt').astype('int') - shift
        #plot(d)
        #plot(ejx[cents[:,0]],d[cents[:,0]],'ro')
        #plot(ejx[cents[:,1]],d[cents[:,1]],'go')
        #show()

        posO, posC = cents[:,0],cents[:,1]
        if mode == 0:
                posO = cents[:,2]
                posC = cents[:,2]
        refO,refC = [],[]
        exap2 = exap + .2*exap
        dev   = exap2/4.
        for i in range(len(posO)):
                if posO[i]-exap2 < 0:
                        x = ejx[:posC[i]+exap2+1]
                        y = d[:posC[i]+exap2+1]
                elif posC[i]+exap2+1 > len(d):
                    x = ejx[posO[i]-exap2:]
                    y = d[posO[i]-exap2:]
                else:   
                        x = ejx[posO[i]-exap2:posC[i]+exap2+1]
                        y = d[posO[i]-exap2:posC[i]+exap2+1]
                y -= y.min()
                tx1 = np.arange(x[0]-dev,x[0],1)
                tx2 = np.arange(x[-1]+1,x[-1]+dev+1,1)
                ty1 = np.zeros(len(tx1))
                ty2 = np.zeros(len(tx2))
                x = np.hstack((tx1,x,tx2))
                y = np.hstack((ty1,y,ty2))

                midi = int(0.5*len(x))
                
                if mode == 1:
                        guess = [np.max(y[:midi]),np.max(y[midi:]),x[0]+np.argmax(y[:midi]),x[0]+midi+np.argmax(y[midi:]),dev,dev]
                        p, success =  scipy.optimize.leastsq(res_gauss2, guess, args=(y,x))
                        refO.append(p[2])
                        refC.append(p[3])
                else:
                        #plot(x,y)
                        #show()
                        #print fcds
                        guess = [np.max(y),x[0]+np.argmax(y),dev]
                        p, success =  scipy.optimize.leastsq(res_gauss, guess, args=(y,x))
                        refO.append(p[1])
                        refC.append(p[1])

        refO = np.array(refO)
        refOt = np.around(refO).astype('int')
        matO = np.zeros((len(refO),sc.shape[1]))
        matO[:,medc] = refO
        refC = np.array(refC)
        refCt = np.around(refC).astype('int')
        matC = np.zeros((len(refC),sc.shape[1]))
        matC[:,medc] = refC

        #plot(d)
        #plot(cents[:,0],d[cents[:,2]],'ro')
        #plot(cents[:,1],d[cents[:,1]],'go')
        #plot(refOt,d[refOt],'ko')
        #plot(refCt,d[refCt],'yo')
        #show()
        #print gfd5
        i = medc -1
        while i >=0:
                d = sc[:,i]
                #if i<1000:
                #       plot(d)
                j = 0
                posO = np.around(refO).astype('int')
                posC = np.around(refC).astype('int')
                trefO,trefC = [],[]
                exap2 = exap + .2*exap
                dev   = exap2/4.

                while j < len(posO):
                        #print j
                        if posO[j]-exap2 < 0:
                                x = ejx[:posC[j]+exap2+1]
                                y = d[:posC[j]+exap2+1]
                        elif posC[j]+exap2+1 > len(d):
                                x = ejx[posO[j]-exap2:]
                                y = d[posO[j]-exap2:]
                        else:   
                                x = ejx[posO[j]-exap2:posC[j]+exap2+1]
                                y = d[posO[j]-exap2:posC[j]+exap2+1]
                        if len(x) < 7:
                                trefO.append(refO[j])
                                trefC.append(refC[j])
                        else:
                                y -= y.min()
                                tx1 = np.arange(x[0]-dev,x[0],1)
                                tx2 = np.arange(x[-1]+1,x[-1]+dev+1,1)
                                ty1 = np.zeros(len(tx1))
                                ty2 = np.zeros(len(tx2))
                                x = np.hstack((tx1,x,tx2))
                                y = np.hstack((ty1,y,ty2))
                                midi = int(0.5*len(x))
                                if mode == 1:
                                        guess = [np.max(y[:midi]),np.max(y[midi:]),x[0]+np.argmax(y[:midi]),x[0]+midi+np.argmax(y[midi:]),dev,dev]
                                        p, success =  scipy.optimize.leastsq(res_gauss2, guess, args=(y,x))
                                        trefO.append(p[2])
                                        trefC.append(p[3])
                                else:
                                        guess = [np.max(y),x[0]+np.argmax(y),dev]
                                        p, success =  scipy.optimize.leastsq(res_gauss, guess, args=(y,x))
                                        trefO.append(p[1])
                                        trefC.append(p[1])

                        j+=1

                orefO = refO.copy()
                trefO = np.array(trefO)
                difO = trefO-refO
                refO = fit_drift(refO,difO,orefO)
                matO[:,i] = refO
                orefC = refC.copy()
                trefC = np.array(trefC)
                difC = trefC-refC
                refC = fit_drift(refC,difC,orefC)
                matC[:,i] = refC

                #if i<500:
                #       plot(d)
                #       I = np.where(refO<2049)[0]
                #       plot(refO[I].astype('int'),d[refO[I].astype('int')],'ko')
                #       show()
                #       print gfd

                i-=4

        refO = matO[:,medc]
        refC = matC[:,medc]

        i = medc +1
        while i < sc.shape[1]:
                d = sc[:,i]
                j = 0
                posO = np.around(refO).astype('int')
                posC = np.around(refC).astype('int')
                trefO,trefC = [],[]
                exap2 = exap + .2*exap
                dev   = exap2/4.

                while j < len(posO):
                        #print j
                        if posO[j]-exap2 < 0:
                                x = ejx[:posC[j]+exap2+1]
                                y = d[:posC[j]+exap2+1]
                        elif posC[j]+exap2+1 > len(d):
                                x = ejx[posO[j]-exap2:]
                                y = d[posO[j]-exap2:]
                        else:   
                                x = ejx[posO[j]-exap2:posC[j]+exap2+1]
                                y = d[posO[j]-exap2:posC[j]+exap2+1]
                        if len(x) < 7:
                                trefO.append(refO[j])
                                trefC.append(refC[j])
                        else:
                                y -= y.min()
                                tx1 = np.arange(x[0]-dev,x[0],1)
                                tx2 = np.arange(x[-1]+1,x[-1]+dev+1,1)
                                ty1 = np.zeros(len(tx1))
                                ty2 = np.zeros(len(tx2))
                                x = np.hstack((tx1,x,tx2))
                                y = np.hstack((ty1,y,ty2))
                                midi = int(0.5*len(x))
                                if mode == 1:
                                        guess = [np.max(y[:midi]),np.max(y[midi:]),x[0]+np.argmax(y[:midi]),x[0]+midi+np.argmax(y[midi:]),dev,dev]
                                        p, success =  scipy.optimize.leastsq(res_gauss2, guess, args=(y,x))
                                        trefO.append(p[2])
                                        trefC.append(p[3])
                                else:
                                        guess = [np.max(y),x[0]+np.argmax(y),dev]
                                        p, success =  scipy.optimize.leastsq(res_gauss, guess, args=(y,x))
                                        trefO.append(p[1])
                                        trefC.append(p[1])
                        j+=1

                orefO = refO.copy()
                trefO = np.array(trefO)
                difO = trefO-refO
                refO = fit_drift(refO,difO,orefO)
                matO[:,i] = refO

                orefC = refC.copy()
                trefC = np.array(trefC)
                difC = trefC-refC
                refC = fit_drift(refC,difC,orefC)
                matC[:,i] = refC

                i+=4

        #imshow(sc_or,vmax=5000)
        for i in range(matO.shape[0]):
                y = matO[i]
                x = np.arange(len(y))
                I = np.where(y!=0)[0]
                x,y = x[I],y[I]
                coef = np.polyfit(x,y,ncoef)
                #plot(x,y,'r')
                residuals = y - np.polyval(coef,x)
                rms = np.sqrt(np.var(residuals))
                I = np.where(np.absolute(residuals)>3*rms)[0]
                cond = True
                if len(I)==0:
                        cond = False
                while cond:
                        im = np.argmax(np.absolute(residuals))
                        x = np.delete(x,im)
                        y = np.delete(y,im)
                        coef = np.polyfit(x,y,ncoef)
                        residuals = y - np.polyval(coef,x)
                        rms = np.sqrt(np.var(residuals))
                        I = np.where(np.absolute(residuals)>3*rms)[0]
                        if len(I)==0:
                                cond = False
                if i == 0:
                        acoefsO = coef
                else:
                        acoefsO = np.vstack((acoefsO,coef))
                #plot(np.polyval(coef,np.arange(len(matO[i]))),'g')
        #show()

        if mode == 1:
                for i in range(matC.shape[0]):
                        y = matC[i]
                        x = np.arange(len(y))
                        I = np.where(y!=0)[0]
                        x,y = x[I],y[I]
                        coef = np.polyfit(x,y,ncoef)
                        #plot(x,y,'r')
                        residuals = y - np.polyval(coef,x)
                        rms = np.sqrt(np.var(residuals))
                        I = np.where(np.absolute(residuals)>3*rms)[0]
                        cond = True
                        if len(I)==0:
                                cond = False
                        while cond:
                                im = np.argmax(np.absolute(residuals))
                                x = np.delete(x,im)
                                y = np.delete(y,im)
                                coef = np.polyfit(x,y,ncoef)
                                residuals = y - np.polyval(coef,x)
                                rms = np.sqrt(np.var(residuals))
                                I = np.where(np.absolute(residuals)>3*rms)[0]
                                if len(I)==0:
                                        cond = False
                        if i == 0:
                                acoefsC = coef
                        else:
                                acoefsC = np.vstack((acoefsC,coef))
                        #plot(np.polyval(coef,np.arange(len(matC[i]))),'b')
                #show()
                #print gfds
                #print acoefs.shape
                return acoefsO, acoefsC
        else:
                return acoefsO
