#!/usr/bin/env python3
# -*- coding: utf-8 -*-

'''
--- utilities.py ---

Just a collection of various functions to be used by the various FORS2 classes,
as well as the __main__ function for running the script.

@place: ESO - Paranal Observatory
@author: Elyar Sedaghati
@year(s):  2018-19
@Telescope(s): UT1
@Instrument(s): FORS2
@Valid for SciOpsPy: v0.1-b
@Documentation url: ***TBD***
@Last SciOps review: 2019-06-16
@Licence: ESO-license or GPLv3 (a copy is given in the root directory of this program)
'''
import numpy as np
from photutils import data_properties
from astropy.stats import sigma_clipped_stats
from photutils import IRAFStarFinder
from astropy.modeling.functional_models import Gaussian2D
from astropy.modeling import fitting
from numpy import array, where, median, abs
from astropy.io import fits
from scipy.optimize import curve_fit
import time, os, pandas

def get_snr(ytemp):
    """get_snr
    
    Estimation of the S/N as described in Stoehr et al. (2008)
    see: https://stdatu.stsci.edu/vodocs/der_snr.pdf)
    The S/N is estimated in 10 ranges (slices) inside the entire 
    spectrum. The max S/N is then taken as the S/N of the spectrum.
    
    Parameter
    ---------
    ytmp
    flux array for which to estimate the S/N.  
    
    Returns
    -------
    maximum S/N among the slices
    """
    n_slices = 10
    a = int(len(ytemp)/10)
    slices_flx = [ytemp[i:i+a] for i in range(0,len(ytemp), a)]
    SNR = []
    for i in slices_flx:
        flux = array(i[where(i != 0.)])
        if len(flux) > 4:
            n = len(flux)
            signal = median(flux)
            noise = 0.6052697 * median(abs(2 * flux[2:n-2] - flux[0:n-4] - flux[4:n] ))
            SNR.append(signal/noise)
        else:
            pass
    SNR = np.sort(SNR)
    return SNR[-2]
###############################################################################
###############################################################################
    
def gaussLine(pix, param):
    """
    gaussian function
    """
    sigma = param['fwhm']
    x = -(pix-param['p0'])**2 / (2*sigma**2)
    return param['offset'] + param['amp']*np.exp(x)

def gauss(x,amp,cen,sigma,offset):
    """
    Returns a gaussian function including an offset.
    """
    return amp*np.exp(-(x-cen)**2/(2*sigma**2))+offset

def lorentz(x, amp, cen, wid, offset):
    """
    Returns a Lorentzian function including an offset.
    """
    return (amp*wid**2/((x-cen)**2+wid**2))+offset

def voigt(x, ampG, cenG, sigmaG, ampL, cenL, widL, offset):
    """
    Returns a Voigt profile including an offset.
    """
    return (ampG*(1/(sigmaG*(np.sqrt(2*np.pi))))*(np.exp(-((x-cenG)**2)/((2*sigmaG)**2))))+\
            ((ampL*widL**2/((x-cenL)**2+widL**2)))
    
###############################################################################
###############################################################################

def get_fwhm(pix, spectrum):
    """get_fwhm
    
    Calculates the fwhm of all peaks for a given array (spectrum), with a 
    peak defined as one that is more than 5 sigma above the noise floor.
    
    ---------------------------------------------------------------------------
    Parameters:
    -----------    
        pix      -- 1D np.array
                    x values of the array where profile/s is/are to be fit.
        spectrum -- 1D np.array
                    y values corresponding to flux where profile/s is/are to be fit.
        
    Returns:
    --------    
        fwhm     -- list
                    Values of fwhm for all the fitted peaks.
        
    ---------------------------------------------------------------------------
    """
    if spectrum.max() > np.median(spectrum)+5*np.std(spectrum):
        fwhm = []
        test = True
        #plt.figure(1)
        #plt.plot(pix,spectrum)
        while test:
            x0    = pix[np.argmax(spectrum)]
            width = 1.5 
            w     = np.where(abs(pix-x0)<30*width)
            
            guess = (spectrum[w].max(),x0,1.5,np.median(spectrum))
            boundaries = ((spectrum[w].max()-0.1*spectrum[w].max(),x0-5,0.3,np.median(spectrum)-20),(spectrum[w].max()+0.1*spectrum[w].max(),x0+5,4.,np.median(spectrum)+20))
            
            pptL, covL = curve_fit(lorentz, pix[w], spectrum[w],guess,bounds=boundaries, maxfev=int(1e7))
            if pptL[2]> 0 and pptL[0]>0:
                fwhm.append(pptL[2])
            spectrum[w] = spectrum[w] - lorentz(pix[w],*pptL) + pptL[-1]
            #plt.plot(pix[w],lorentz(pix[w],*pptL),c='red',alpha=0.5)
            test = spectrum.max() > np.median(spectrum)+5*np.std(spectrum)
        #plt.plot(pix,spectrum)
        #plt.show()
        return fwhm
###############################################################################
###############################################################################
        
def get_IMG_IQ_ELL(infile):
    """get_IMG_IQ_ELL
    
    This is function to calculate the image quality and ellipticity for a 
    FORS2 image, taken in IMG mode.
    ---------------------------------------------------------------------------
    Parameters:
    -----------
        infile    -- .fits 
                    input FORS2 IMG/IPOL fits file
        
    Returns:
    --------
        IQ,ELL    -- float,float
                     IQ  = median(fwhm) * pix_scale * binning factor
                     ell = median(ell)
    ---------------------------------------------------------------------------
    Functionality description:
        1) The edges of the image are cut to avoid artefacts.
        2) IRAFStarFinder algorithm is used to identify non-saturated 
           point-sources.
        3) The brightest 40% of these are selected for analysis.
        4) Sub-arrays of 20x20 are cut around each star, if the centre of the
           source is not close to any edge.
        5) These arrays are median subtracted.
        6) A 2D Gaussian model is fitted to the matrix using a Levensburg-
           Marquardt least-squares fitter
        7) fwhm of each source is calculated as the mean of the fwhm of the 
           2D profile along both x and y directions.
        8) IQ is then calculated as the median of these values for those
           sources whose ellipticity is below 0.3, times the pixel scale, times
           the binning factor.
        9) The ellipticity is the median of the ellpticity of those source, who
           who have this value below 0.3.
    ---------------------------------------------------------------------------

    """
    hdu = fits.open(infile)
    pix_scale = hdu[0].header['ESO INS PIXSCALE']
    binning = hdu[0].header['ESO DET WIN1 BINY']
    asm_seeing = hdu[0].header['ESO TEL IA FWHM']
    data = hdu[0].data[int(10/binning):int(1890/binning),int(390./binning):int(3670/binning)]
    data = data.astype(float)
    mean, median, std = sigma_clipped_stats(data, sigma=3., maxiters=5)
    iraffind = IRAFStarFinder(fwhm=1.5*asm_seeing/(pix_scale*binning), threshold=3.*std, exclude_border=True)
    sources = iraffind(data)
    sources.sort('flux')
    sources = sources[-1*int(0.4*len(sources)):]
    
    ell = []; FWHM = []        
    cols = len(data[0]); rows = len(data)

    for x,y,z in zip(sources['xcentroid'],sources['ycentroid'],sources['flux']):
        if 50<x<0.9*cols-10 and 30<y<0.95*rows-10 :
            data2 = data[int(y-10):int(y+10),int(x-10):int(x+10)]
            mean, median, std = sigma_clipped_stats(data2, sigma=3., maxiters=5)
            data2 = data2 - median
            data2 = data2.astype(float)
            
            gauss_2d = Gaussian2D(amplitude=data2[10][10], x_mean=10, y_mean=10, x_stddev=None, y_stddev=None, theta=None)
            fitter = fitting.LevMarLSQFitter()
            l,h = np.mgrid[:20,:20]
            p = fitter(gauss_2d, h, l, data2)
            fwhm = np.mean([p.x_fwhm,p.y_fwhm])
            cat = data_properties(data2)
            if cat.ellipticity.value < 0.3:
                ell.append(cat.ellipticity.value)
                FWHM.append(fwhm)
    del data
    return np.median(FWHM)*pix_scale*binning, np.median(ell)
###############################################################################
###############################################################################

def get_IPOL_IQ_ELL(infile):
    """get_IPOL_IQ_ELL
    
    This is function to calculate the image quality and ellipticity for a 
    FORS2 image, taken in IPOL mode.
    ---------------------------------------------------------------------------
    Parameters:
    -----------
        infile    -- .fits 
                    input FORS2 IMG/IPOL fits file
        
    Returns:
    --------
        IQ,ELL    -- float,float
                     IQ  = median(fwhm) * pix_scale * binning factor
                     ell = median(ell)
    ---------------------------------------------------------------------------
    Functionality description:
        1) The edges of the image are cut to avoid artefacts.
        2) IRAFStarFinder algorithm is used to identify non-saturated 
           point-sources.
        3) The brightest 40% of these are selected for analysis.
        4) Sub-arrays of 20x20 are cut around each star, if the centre of the
           source is not close to any edge.
        5) These arrays are median subtracted.
        6) A 2D Gaussian model is fitted to the matrix using a Levensburg-
           Marquardt least-squares fitter
        7) fwhm of each source is calculated as the mean of the fwhm of the 
           2D profile along both x and y directions.
        8) IQ is then calculated as the median of these values for those
           sources whose ellipticity is below 0.3, times the pixel scale, times
           the binning factor.
        9) The ellipticity is the median of the ellpticity of those source, who
           who have this value below 0.3.
    ---------------------------------------------------------------------------

    """

    hdu = fits.open(infile)
    header = hdu[0].header
    pix_scale = header['ESO INS PIXSCALE']
    binning = header['ESO DET WIN1 BINX']
    asm_seeing = header['ESO TEL IA FWHM']
    
    rows = len(hdu[0].data)
    cols = len(hdu[0].data[0])
    data = hdu[0].data[int(0.03*rows):int(0.89*rows),int(0.3*cols):int(0.7*cols)]
    mean, median, std = sigma_clipped_stats(data, sigma=3., maxiters=5)
    data = data.astype(float)
    daofind = IRAFStarFinder(fwhm=1.5*asm_seeing/(pix_scale*binning), threshold=3.*std, exclude_border=True)
    sources = daofind(data)
    sources.sort('flux')
    sources = sources[-1*int(0.4*len(sources)):]
    ell = []; FWHM = []

    if header['ESO TEL TRAK STATUS'] == 'NORMAL': #Normal observing modes.
        cols = len(data[0]);rows = len(data) 
        for x,y,z in zip(sources['xcentroid'],sources['ycentroid'],sources['flux']):
            if 20<x<0.95*cols-10 and 20<y<0.95*rows-10 :
                data2 = data[int(y-10):int(y+10),int(x-10):int(x+10)]
                mean, median, std = sigma_clipped_stats(data2, maxsigma=3., maxiters=5)
                data2 = data2 - mean
                data2 = data2.astype(float)
            
                gauss_2d = Gaussian2D(amplitude=data2[10][10], x_mean=10, y_mean=10, x_stddev=None, y_stddev=None, theta=None)
                fitter = fitting.LevMarLSQFitter()
                l,h = np.mgrid[:20,:20]
                p = fitter(gauss_2d, h, l, data2)
                fwhm = np.mean([p.x_fwhm,p.y_fwhm])
                cat = data_properties(data2)
                if cat.ellipticity.value < 0.3:
                    ell.append(cat.ellipticity.value)
                    FWHM.append(fwhm)
        return np.median(FWHM)*pix_scale*binning,np.median(ell)

    else:                 #Moving targets
        os.system('''wget -O temp.csv "http://archive.eso.org/wdb/wdb/asm/dimm_paranal/query?start_date=%s..%s&tab_airmass=1&tab_rfl=0&tab_rfl_time=0&wdbo=csv" -q'''%(header['DATE-OBS'][:10]+' '+header['DATE-OBS'][11:],header['DATE'][:10]+' '+header['DATE'][11:]))
        df = pandas.read_csv('temp.csv')
        if df.columns.values[0] == '# No data returned !':
            try : see = header['ESO TEL IA FWHM']
            except KeyError: see = 0.0
        else:
            try:
                seeing_dimm = np.asarray(df['DIMM Seeing ["]'][:-5])
                airmass_dimm = np.asarray(df['Airmass'][:-5])
                airmass = np.mean([header['ESO TEL AIRM START'],header['ESO TEL AIRM END']])
                seeing = []
                for i,j in zip(seeing_dimm,airmass_dimm):
                    seeing.append(i * (airmass/j)**0.6)
                see = np.mean(seeing)
            except KeyError: see = 0.0
            os.system('rm temp.csv')
        ell = 0.
        return see,ell