#!/usr/bin/env python3
# -*- coding: utf-8 -*-
'''
--- FORS2_QC.py ---

Main class for the determination of FORS2 QC parameters, except IPOL & PMOS.

Header &  QC parameter definitions: 
==============================================================
OB name: 'ESO OBS NAME' / OB name                             
Prog ID: 'ESO OBS PROG ID' / ESO program identification         
OB ID: 'ESO OBS ID' / Observation block ID                  
Date:  'DATE' / file creation date (YYYY-MM-DDThh:mm:ss UT)
Requested transparency: 'ESO OBS AMBI TRANS' / Req. sky transparency
Airmass max: max('ESO TEL AIRM START','ESO TEL AIRM END')
Airmass req: 'ESO OBS AIRM' / Req. max. airmass 
Airmass per: computed / (airmass_max - airmass_req)/airmass_req
Airmass grade: < 0% --> A | <= 10% --> B | > 10% --> C                
SEEING: if IMG mode: 
            'ESO QC IMGQU' / Image quality of scientific exposure for IMG mode only
            (although this seems to be wrong and is to be calculated directly on the reduced images)
        elif IPOL mode:
            FWHM * pix_scale * binning / FWHM measured on 40% brightest point sources
        elif LSS or PMOS modes:
            FWHM * pix_scale * binning / FWHM measured on the spectral trace
        elif MOS or MXU mode:
            DIMM seeing corrected for airmass and adjusted for image analysis of ACQ frames.
SEE_REQ: 'ESO OBS AMBI FWHM' / Req. max. seeing
SEE_PER: '(seeing - see_requested)/(see_requested)' for IMG mode only.
Seeing grade: < 0% --> A | <= 10% --> B | > 10% --> C for IMG mode only
SEE_PER: % of DIMM readings during observations out of range. (all other modes)
Seeing_grade: <= 10% A | <= 20% --> B | > 20% --> C

Ellipticity : if IMG mode:
                'ESO QC IMGQUELL' / Mean star ellipticity
                (again the pipeline value seems wrong and this shall be measured durectly on the images)
              elif IPOL mode:
                from 2DGaussian fit to the 40% brightest point sournces
              else:
                  0.
Ell. grade: <= 0.10 --> A | <= 0.20 --> B | > 0.20 --> C 
'''

from astropy.io import fits
import os, pandas
import numpy as np
from utilities import get_snr, get_fwhm, get_IMG_IQ_ELL, get_IPOL_IQ_ELL
from raw_IPOL_QC import IPOL_QC
from PMOS_QC import PMOS_QC, PMOS2_QC
from hardcoded import *

class FORS2_QC (object) :
    """
    FORS2_QC
    ========
    
    Main class for the determination of FORS2 QC grade.  It works for all
    observation modes, except for PMOS and IPOL which have their own classes.
    
    Inputs:
        -- reduced FORS2 fits frame
        -- raw FORS2 fits frame
        
    Output:
        Prints out a series of file information values, QC evaluations & a 
        final suggested grade.
        
    Methods:
        
        image_grade(): calculates grades for various QC parameters.
        
        qc_onscreen(): prints out the results to the terminal.
        
        
    --------------------------------------------------------------------------    
    """
    def __init__(self, FORS2red, FORS2acq, FORS2raw):
        
        self.hdu = fits.open(FORS2red)
        self.hdr = self.hdu[0].header
        self.hdu2 = fits.open(FORS2raw)
        self.hdr2 = self.hdu2[0].header
        self.pix_scale = self.hdr2['ESO INS PIXSCALE']
        self.binning = self.hdr2['ESO DET WIN1 BINY']
        try: self.am_max = max(self.hdr['ESO TEL AIRM START'],self.hdr['ESO TEL AIRM END'])
        except KeyError: self.am_max = 1.0
        self.am_req = self.hdr['ESO OBS AIRM']
        self.am_per = None
        self.am_grade = None
        ###########################################################
        ##### Image Quality & Ellipticity Calculations ############
        ###########################################################
        if self.hdr['ESO INS MODE'] == 'IMG' and self.hdr['ESO PRO CATG'] == 'SCIENCE_REDUCED_IMG':
            if self.hdr['ESO TEL TRAK STATUS'] == 'NORMAL':
                try:
                    self.see = get_IMG_IQ_ELL(FORS2red)[0]
                    self.ell = get_IMG_IQ_ELL(FORS2red)[1]
                except MemoryError:
                    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'''%(self.hdr['DATE-OBS'][:10]+' '+self.hdr['DATE-OBS'][11:],self.hdr['DATE'][:10]+' '+self.hdr['DATE'][11:]))
                    df = pandas.read_csv('temp.csv')
                    # try doing this with JSON to read data directly
                    if df.columns.values[0] == '# No data returned !':
                        try : self.see = self.hdr2['ESO TEL IA FWHM']
                        except KeyError: self.see = 0.0
                    else:
                        try:
                            seeing_dimm = np.asarray(df['DIMM Seeing ["]'][:-5])
                            airmass_dimm = np.asarray(df['Airmass'][:-5])
                            airmass = np.mean([self.hdr['ESO TEL AIRM START'],self.hdr['ESO TEL AIRM END']])
                            seeing = []
                            for i,j in zip(seeing_dimm,airmass_dimm):
                                seeing.append(i * (airmass/j)**0.6)
                            self.see = np.mean(seeing)
                        except KeyError: self.see = 0.0
                    os.system('rm temp.csv')
                    self.ell = 0.

            else:
                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'''%(self.hdr['DATE-OBS'][:10]+' '+self.hdr['DATE-OBS'][11:],self.hdr['DATE'][:10]+' '+self.hdr['DATE'][11:]))
                df = pandas.read_csv('temp.csv')
                # try doing this with JSON to read data directly
                if df.columns.values[0] == '# No data returned !':
                    try : self.see = self.hdr2['ESO TEL IA FWHM']
                    except KeyError: self.see = 0.0
                else:
                    try:
                        seeing_dimm = np.asarray(df['DIMM Seeing ["]'][:-5])
                        airmass_dimm = np.asarray(df['Airmass'][:-5])
                        airmass = np.mean([self.hdr['ESO TEL AIRM START'],self.hdr['ESO TEL AIRM END']])
                        seeing = []
                        for i,j in zip(seeing_dimm,airmass_dimm):
                            seeing.append(i * (airmass/j)**0.6)
                        self.see = np.mean(seeing)
                    except KeyError: self.see = 0.0
                os.system('rm temp.csv')
                self.ell = 0.

        elif self.hdr['ESO PRO CATG'] in ['REDUCED_SCI_LSS','REDUCED_SCI_PMOS'] and self.hdr['ESO DET CHIP1 ID']=='CCID20-14-5-3':
            DIMM = self.hdu[0].header['ESO TEL IA FWHM']
            temp = []
            for infile in FORS2acq:
                temp.append(get_IMG_IQ_ELL(infile)[1])
                #temp.append(-99) <--RTh in case the previous line crashes
            self.ell = np.mean(temp)
            try:
                a = int(np.shape(self.hdu2[0].data[0])[0])
                temp = []
                for l in np.arange(a/2,a,20):
                    data = self.hdu2[0].data[int(60/self.binning):int(1900/self.binning),int(l)]
                    data = data.flatten()
                    x = np.arange(len(data))
                    res = get_fwhm(x, data)*self.pix_scale*self.binning
                    if res:
                        for i in res:
                            if 0.7*DIMM < i < 1.3*DIMM:
                                temp.append(i)
                self.see = np.median(temp)
            except TypeError:
                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'''%(self.hdr['DATE-OBS'][:10]+' '+self.hdr['DATE-OBS'][11:],self.hdr['DATE'][:10]+' '+self.hdr['DATE'][11:]))
                df = pandas.read_csv('temp.csv')
                if df.columns.values[0] == '# No data returned !':
                    try : self.see = self.hdr2['ESO TEL IA FWHM']
                    except KeyError: self.see = 0.0
                else:
                    try:
                        seeing_dimm = np.asarray(df['DIMM Seeing ["]'][:-5])
                        airmass_dimm = np.asarray(df['Airmass'][:-5])
                        airmass = np.mean([self.hdr['ESO TEL AIRM START'],self.hdr['ESO TEL AIRM END']])
                        seeing = []
                        for i,j in zip(seeing_dimm,airmass_dimm):
                            seeing.append(i * (airmass/j)**0.6)
                        #self.see = "{0:.2f}".format(np.mean(seeing))
                        self.see = np.mean(seeing)
                    except KeyError: self.see = 0.0
                os.system('rm temp.csv')
            else:
                self.see = np.median(temp)
#            self.ell = 0.
        elif self.hdr['ESO PRO CATG'] in ['REDUCED_SCI_MXU', 'REDUCED_SCI_MOS'] and self.hdr['ESO DET CHIP1 ID']=='CCID20-14-5-3':
            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'''%(self.hdr['DATE-OBS'][:10]+' '+self.hdr['DATE-OBS'][11:],self.hdr['DATE'][:10]+' '+self.hdr['DATE'][11:]))
            df = pandas.read_csv('temp.csv')
            if df.columns.values[0] == '# No data returned !':
                try : self.see = self.hdr2['ESO TEL IA FWHM']
                except KeyError: self.see = 0.0
            else:
                try:
                    seeing_dimm = np.asarray(df['DIMM Seeing ["]'][:-5])
                    airmass_dimm = np.asarray(df['Airmass'][:-5])
                    airmass = np.mean([self.hdr['ESO TEL AIRM START'],self.hdr['ESO TEL AIRM END']])
                    seeing = []
                    for i,j in zip(seeing_dimm,airmass_dimm):
                        seeing.append(i * (airmass/j)**0.6)
                    self.see = np.mean(seeing)
                except KeyError: self.see = 0.0
            os.system('rm temp.csv')
            temp = []
            for infile in FORS2acq:
                temp.append(get_IMG_IQ_ELL(infile)[1])
            self.ell = np.mean(temp)
        ###########################################################
        ############ Signal to Noise Calculations #################
        ###########################################################
        if self.hdr['ESO PRO CATG'] == 'REDUCED_SCI_LSS':
            temp = []
            for i in np.arange(np.shape(self.hdu[0].data)[0]):
                temp.append(get_snr(self.hdu[0].data[i]))
            self.snr = str("{0:.2f}".format(max(temp)))
        elif self.hdr['ESO PRO CATG'] == 'REDUCED_SCI_PMOS':
            for i in range(1,len(self.hdu)):
                temp = []
                for j in np.arange(np.shape(self.hdu[i].data)[0]):
                    temp.append(get_snr(self.hdu[i].data[j]))
            self.snr = str("{0:.2f}".format(np.max(temp)))
        elif self.hdr['ESO PRO CATG'] == 'REDUCED_SCI_MXU':
            try:
                temp = []
                for i in np.arange(np.shape(self.hdu[0].data)[0]):
                    temp.append(get_snr(self.hdu[0].data[i]))
                temp = np.sort(temp)
                self.snr = str("{0:.2f}".format(temp[-2])) #SNR:
            except IndexError:
                self.snr = 'N/A'
        elif self.hdr['ESO PRO CATG'] == 'REDUCED_SCI_MOS':
            temp = []
            for i in np.arange(np.shape(self.hdu[0].data)[0]):
                temp.append(get_snr(self.hdu[0].data[i]))
            self.snr = str("{0:.2f}".format(max(temp)))
            #self.snr = np.max(temp) #SNR    
        else: self.snr = 'N/A'
        ###########################################################
        ################## QC attributes ##########################
        ###########################################################
        self.see_req = self.hdr['ESO OBS AMBI FWHM']
        self.see_spec = None
        self.see_per = None
        self.see_grade = None
        self.ell_grade = None
        self.sugg_grade = None

    def image_grade (self) :            
        """
        Calculates grades based on various metrics. 
        """
        # Airmass
        # Compute airmass percentage as (airmass - airmass_requested)/(airmass_requested)
        self.am_per = (self.am_max-self.am_req)/self.am_req
        if self.am_per < 0 : self.am_grade = 'A'
        elif self.am_per <= 0.1 : self.am_grade = 'B'
        else : self.am_grade = 'C'
        
        #Seeing        
        self.see_per = (self.see-self.see_req)/self.see_req
        if self.see_per < 0 : self.see_grade = 'A'
        elif self.see_per <= 0.1 : self.see_grade = 'B'
        else : self.see_grade = 'C'
        
        #Ellipticity
        if self.ell == 0.: self.ell_grade = 'N/A'
        elif self.ell <= 0.10 : self.ell_grade = 'A'
        elif self.ell <= 0.20 : self.ell_grade = 'B'
        else : self.ell_grade = 'C'
        
        if self.ell_grade == 'N/A': self.sugg_grade = max(self.am_grade, self.see_grade)
        else: self.sugg_grade = max(self.am_grade, self.see_grade, self.ell_grade)


    # Method to print on screen the quality control results and to save them to a txt file     
    def qc_onscreen (self) :
        if self.sugg_grade == 'A': 
            CSTART = '\033[32m'
            CEND   = '\033[0m'
        elif self.sugg_grade == 'B':
            CSTART = '\033[33m'
            CEND   = '\033[0m'
        else:
            CSTART = '\033[31m'
            CEND   = '\033[0m'
        mystring = '{:<23s} {:4s}  {:>15s} {:<7s} {:<20s} {:<4s} | {:1.2f}  {:4.2f} {:>2s} | {:3.3f} {:>4s} | {:4.2f}  {:4.2f}  {:>2s} | {:>6s}  |   {:2s}'.format(
                self.hdr['ESO OBS NAME'][:23], self.hdr['ESO INS MODE'],self.hdr['ESO OBS PROG ID'] , str(self.hdr['ESO OBS ID']), 
                self.hdr['DATE'], self.hdr['ESO OBS AMBI TRANS'], self.am_max, self.am_req, 
                self.am_grade, self.ell, self.ell_grade, np.median(self.see), self.see_req,
                self.see_grade, self.snr ,CSTART+self.sugg_grade+CEND)
        print (mystring)