
import pickle

import numpy as np
from scipy import stats, signal

from astropy.io import fits
from astropy import modeling as md

from PySpectrograph.Spectra import Spectrum

from GiraffeModel import GiraffeModel

def get_wavelength(xarr, order, gamma=6.3, cam_foc=475, xpos=0.0):
    """Based on the model for the spectragraph, return the wavelength
       at a given x-position.

    Parameters
    ----------
    xarr: numpy.ndarray
        Array of x-positions
 
    order: int
        Order of the spectra

    gamma: float
        Angle between the grating and camera
  
    cam_foc: float
        Camera focus

    xpos: float
        X-position of the detector relative to the central ray

    Returns
    -------
    warr: numpy.ndarray
        Array of wavelengths in Angstroms corresponding to xarr
 

    """
    g = GiraffeModel(order=order, gamma=gamma)
    xc=g.detector.get_xpixcenter()
    g.camera.focallength = cam_foc
    g.detector.xpos = xpos
    warr = 1e7* g.get_wavelength(xarr, gamma=gamma) 
    return warr
 

def get_data(data, coef, dy=5):
    """Based on the shape given by coef with a width of dy, extract
       the data associated with the shape

    Parameters
    ----------
    data: numpy.ndarray
        A 2-D array to extract data

    coef: numpy.ndarray
        A set of coefficients to describe the shape of an order.  These 
        coefficients should be generated by numpy.polyfit.

    dy: int
        Width of area around the line describing the order that should
        be extracted.

    Returns
    -------
    cdata: numpy.ndarray
        A 1-D array containing the flux from the extracted order

    """
    yi, xi = np.indices(data.shape)
    xarr  = np.arange(len(data[0])) 
    yarr  = np.polyval(coef, xarr)
    mask = (yi > yarr-dy) * (yi < yarr+dy)
    cdata = 1.0 * data
    cdata[mask==0] = 0
    return cdata.sum(axis=0)

def getspec(arcfile='ThAr.salt', y=None, dw=0.01, res=0.1):
    """Read in an arc file from a txt file

    Returns 
    -------
    spec: ~PySpectrograph.Spectrum.Spectrum
        Spectrum object with the lines for the Arc lamp
    """
    if y is None:
       wrange = None
    else:
       wrange = [y.min(), y.max()]
    sw,sf = np.loadtxt('data/ThAr.salt', usecols=(0,1), unpack=True)
    spec=Spectrum.Spectrum(sw, sf, wrange=wrange, dw=dw, stype='line', sigma=res)
    return spec

def getfitsspec(arcfile='thar.fits', dw=0.01, res=0.1):
    """Read in an arc file from a FITS file

    Returns 
    -------
    spec: ~PySpectrograph.Spectrum.Spectrum
        Spectrum object with the lines for the Arc lamp
    """
    wrange = None 
    shdu = fits.open(arcfile)
    ctype1=shdu[0].header['CTYPE1']
    crval1=shdu[0].header['CRVAL1']
    cdelt1=shdu[0].header['CDELT1']
    sf=shdu[0].data
    sw=crval1+cdelt1*np.arange(len(shdu[0].data))
    spec=Spectrum.Spectrum(sw, sf, wrange=wrange, dw=dw, stype='continuum', sigma=res)
    return spec

def ncor(x, y):
    """Calculate the normalized correlation of two arrays"""
    d=np.correlate(x,x)*np.correlate(y,y)
    if d<=0: return 0
    return np.correlate(x,y)/d**0.5


def extract_data(data, orders, ys):
   """
 
   Parameters
   ----------
   data: numpy.ndarray
        A 2-D array to extract data

   orders: numpy.ndarray
        An array containing the y-position and coefficients
        describing the shape of each orders

   ys: int
        Y-position of the order to extract


   Returns
   -------
   xarr: numpy.ndarray
        A 1-D array describing the x-position of the flux

   farr: numpy.ndarray
        A 1-D array describing flux fo the extracted order
   
   """
   i = abs(orders[:,0]-ys).argmin()
   farr  = get_data(data, orders[i,1:])
   xarr = np.arange(len(farr))
   return xarr, farr


def match_order(xarr, farr, n1, n2, arcfile='data/thar.fits', dw=20, gamma=6.28, cam_foc=475):
    """Find the order that best matches the extract data.  

   
    """
 
    spec = getfitsspec(arcfile=arcfile)

    orders_arr = np.arange(n1,n2)
    bestcc_arr = 0.0*orders_arr
    for j,n in enumerate(orders_arr):
	y = get_wavelength(xarr, n, gamma = gamma, cam_foc=cam_foc)
        mask = (spec.wavelength > y.min()-dw) * (spec.wavelength < y.max()+dw)
        sw = spec.wavelength[mask]
        sf = spec.flux[mask]
        nsf = np.interp(y, sw, sf)
        
        #allow for a zeropoint shift
        dx_arr = np.arange(-100,100,2)
        cc_arr = 0.0*dx_arr
        for i in range(len(dx_arr)):
            xnsf = np.interp(xarr+dx_arr[i], xarr, nsf)
            cc_arr[i] = ncor(farr, xnsf)
            
        bestcc_arr[j] = cc_arr.max()
    n = orders_arr[bestcc_arr.argmax()]
    return n

def find_first_order(data, orders, arcfile='data/thar.fits', n1=60, n2=160, dw=10, gamma=6.28, cam_foc=475):
    """Find the first order in the array
    """
    ypos = orders[:,0]
    ord_arr = 0.0 * ypos
    for i in range(len(ypos)):
        yc = ypos[i]
        xarr, farr = extract_data(data, orders, yc)
        norder = match_order(xarr, farr, n1, n2, arcfile=arcfile, dw=dw, gamma=gamma, cam_foc=cam_foc)
        ord_arr[i] = norder

    xarr = np.arange(len(ord_arr))
    return stats.mode(ord_arr-xarr)[0][0]


def fit_gamma(xarr, f, n, arcfile='thar.fits', dw=10, gamma=6.413, dg1=-0.3, dg2=0.3, dgs=0.01,
              cam_foc=475):
    """Determine the best value for gamma
    """
    #set up the data
    y = get_wavelength(xarr, n, gamma = gamma, cam_foc=cam_foc)
    spec = getfitsspec(arcfile=arcfile)
    mask = (spec.wavelength > y.min()-dw)*(spec.wavelength < y.max()+dw)
    if mask.sum()==0: return gamma
    sw = spec.wavelength[mask]
    sf = spec.flux[mask]
    gamma_arr = np.arange(dg1, dg2, dgs)
    cc_arr = 0.0*gamma_arr
    for i, dg in enumerate(gamma_arr):
        y = get_wavelength(xarr, n, gamma = gamma+dg, cam_foc=cam_foc)
        nsf = np.interp(y, sw, sf)
        cc_arr[i] = ncor(f, nsf)
    dg=gamma_arr[cc_arr.argmax()]
    y = get_wavelength(xarr, n, gamma = gamma+dg, cam_foc=cam_foc)
    return gamma+dg

def fit_order(xarr, f, sw, sf, ws,  dw=3, gamma=6.413, cam_foc=475):
    """Fit a spectrograph model to the data
 
    """

    fit_g = md.fitting.LevMarLSQFitter()

    xp  = signal.find_peaks_cwt(f, np.array([5]))
    xp = np.array(xp)
    fp = f[xp]
    wp = ws(xp)
    y = ws(xarr)

    #find the best matches
    mx = []
    mw = []
    for i in fp.argsort()[::-1][0:]:
        mask = (y>wp[i] - dw) * ( y< wp[i]+dw)
        g = md.models.Gaussian1D(amplitude=f[mask].max(), mean=xp[i], stddev=0.5)
        g = fit_g(g, xarr[mask], f[mask])
        x=g.mean.value
        s=3*g.stddev.value
        smask = (sw < ws(x-s))*(sw > ws(x+s))
        if smask.sum() == 1:
           mx.append(x)
           mw.append(sw[smask][0])
        else:
           pass

    return mx, mw

def wavelength_solution(xarr, farr, n, arcfile='data/thar_list.txt', 
                        func_order=2, gamma=6.413, cam_foc=474.2, niter=3):
    """Determine the wavelength solution for the data
    """
    y = get_wavelength(xarr, n, gamma = gamma, cam_foc=cam_foc)
    sw = np.loadtxt(arcfile, usecols=(0,), unpack=True)
    sf = 0.0*sw+1

    fit_g = md.fitting.LevMarLSQFitter()
    ws_init = md.models.Legendre1D(func_order)
    ws = fit_g(ws_init, xarr, y)

    for i in range(niter):
       xp, wp = fit_order(xarr, farr, sw, sf, ws, dw=3, gamma=gamma, 
                          cam_foc=cam_foc)
       try:
           ws = fit_g(ws_init, xp, wp)
       except:
           return ws
    return ws

def fit_all_orders(data, orders, n1, outfile, arcfile='data/thar.fits', wsarcfile='data/thar_list.txt',
                   dw=3, gamma = 6.4, cam_foc=475):
    ypos = orders[:,0]
    sol_dict={}
    for i, yc in enumerate(ypos):
        xarr, farr = extract_data(data, orders, yc)
        gamma = fit_gamma(xarr, farr, n1+i, arcfile=arcfile, dw=dw, gamma=gamma, cam_foc=cam_foc)
        ws = wavelength_solution(xarr, farr,  n1+i, arcfile=wsarcfile, gamma = gamma, cam_foc=cam_foc)
        sol_dict[n1+i] = [yc, gamma, cam_foc, orders[i], ws]

    pout = open(outfile, 'wb')
    pickle.dump(sol_dict, pout)
    return sol_dict

def extract_spectra(ccd, sol_dict, dy=5, outfile=None):

    keys = sol_dict.keys()
    keys.sort()
    all_xarr = None
    for n in keys:
        yc, gamma, cam_foc, orders, ws = sol_dict[n]
        # extract the orders
        farr  = get_data(ccd.data, orders[1:], dy=dy)
        # apply wavelength solution
        xarr = np.arange(len(farr))
        warr = ws(xarr)
        if all_xarr is None:
           all_xarr = 1.0 * xarr
           all_warr = 1.0 * warr
           all_farr = 1.0 * farr
           all_farr_err = 1.0 * abs(farr)**0.5
           all_order = 0.0 * xarr + n
        else:
           all_xarr = np.concatenate((all_xarr, xarr))
           all_warr = np.concatenate((all_warr, warr))
           all_farr = np.concatenate((all_farr, farr))
           all_farr_err = np.concatenate((all_farr_err, 1.0*abs(farr)**0.5))
           all_order = np.concatenate((all_order, 0.0 * xarr + n))

    #estimate the signal to noise
    sn = (abs(all_farr)**0.5).mean()
    print sn
    if sn<0: sn = 0

    if outfile is not None:
        colx = fits.Column(name='x', format='D', array=all_xarr)
        colw = fits.Column(name='wavelength', format='E', array=all_warr)
        colf = fits.Column(name='flux', format='E', array=all_farr)
        cole = fits.Column(name='flux_err', format='E', array=all_farr_err)
        colo = fits.Column(name='order', format='E', array=all_order)
        cols = [colx, colw, colf, cole, colo]

        tbhdu = fits.BinTableHDU.from_columns(cols)
        tbhdu.header.update('SN', sn, 'Mean Signal to Noise')
        tbhdu.writeto(outfile, clobber=True) 
    return all_xarr, all_warr, all_farr, all_farr_err, all_order, sn



