import reflex 
import sys 
import os
#from astropy.io import fits
import pyfits as fits
import numpy as np
import time

def getw(header):
   N=header["NAXIS1"]
   CRPIX=header["CRPIX1"]
   CDELT=header["CDELT1"]
   CRVAL=header["CRVAL1"]
   starty=-(CRPIX-1.0)*CDELT+CRVAL
   w=np.arange(N)*CDELT + starty
   return w


if __name__ == '__main__':
 #getting the inputs and defining the outputs
  parser = reflex.ReflexIOParser()
  parser.add_option("-i", "--in_sof", dest="in_sof")
  parser.add_output("-o", "--out_sof", dest="out_sof")
  parser.add_input("-p", "--in_sop", dest="in_sop", type='string')
  parser.add_output("-m", "--messages", dest="messages")

  inputs  = parser.get_inputs()
  outputs = parser.get_outputs()
  
  in_sof = inputs.in_sof
  pattern = '--products-dir'
  for arg in sys.argv:
     if arg.split("=")[0] == pattern:
         output_dir = arg.split("=")[1]

  files=in_sof.files
  params  = inputs.in_sop
  values= {}
  for param in params:
          values[param.name]= param.value

  strategy = values['define_correction.telluric']
  if strategy > 3 or strategy < 0:
        strategy = 0
        outputs.messages= 'Warning: Input parameters must be integer within 0 and 3. Setting it to 0.'

         
  #read response if present
  IsResponsePresent = 0
  for file in files:
   if file.category == 'RESPONSE':
      hdu_res = fits.open(file.name)
      IsResponsePresent = 1 
      name_response = file.name
      output_purpose = file.purposes
  #read telluric if present
  IsTelluricPresent = 0
  for file in files:
   if file.category == 'TELLURIC':
      hdu_telres = fits.open(file.name)
      IsTelluricPresent = 1
      name_telluric = file.name
      output_purpose = file.purposes

  #read telluric only if present
  IsTelluricOnlyPresent = 0
  for file in files:
   if file.category == 'TELLURIC_ONLY':
      hdu_tel = fits.open(file.name)
      IsTelluricOnlyPresent = 1
      name_telluric_only = file.name
      output_purpose = file.purposes

#  print   'IsTelluricPresent', IsTelluricPresent
#  print   'IsTelluricOnlyPresent', IsTelluricOnlyPresent
#  print   'IsResponsePresent',IsResponsePresent
#  print 'strategy',strategy
  #GET THE ZEROPOINTS FROM THE TELLURIC FILE, if present and not used.
  if IsTelluricPresent == 1 and strategy !=0:
          hdu_tell = fits.open(name_telluric)
          #compute mean zpt and thresholds
          zpt_1 = list()
          thr_1 = list()
          zpt_2 = list()
          thr_2 = list()
          zpt_3 = list()
          thr_3 = list()
         
          for exten in range(1,17,2):
            data = hdu_tell[exten].data
            if data != None:
              zpt_1.append(hdu_tell[exten].header["HIERARCH ESO QC ZPOINT"])
              thr_1.append(hdu_tell[exten].header["HIERARCH ESO QC THRUPUT"])

          #compute mean zpt for chip2
          for exten in range(17,33,2):
            data = hdu_tell[exten].data
            if data != None:
              zpt_2.append(hdu_tell[exten].header["HIERARCH ESO QC ZPOINT"])
              thr_2.append(hdu_tell[exten].header["HIERARCH ESO QC THRUPUT"])

          #compute mean zpt for chip3
          for exten in range(33,49,2):
            data = hdu_tell[exten].data
            if data != None:
              
              zpt_3.append(hdu_tell[exten].header["HIERARCH ESO QC ZPOINT"])
              thr_3.append(hdu_tell[exten].header["HIERARCH ESO QC THRUPUT"])
        
          zpt_1m = np.mean(zpt_1)
          thr_1m = np.mean(thr_1)
          zpt_2m = np.mean(zpt_2)
          thr_2m = np.mean(thr_2)
          zpt_3m = np.mean(zpt_3)
          thr_3m = np.mean(thr_3)

       
#REDEFINE THE ZPT in the header. Updated header is used but not saved
          for exten in range(1,17,2):
            data = hdu_tell[exten].data
            if data == None:
              hdu_tell[exten].header["HIERARCH ESO QC ZPOINT"] = zpt_1m
              hdu_tell[exten].header["HIERARCH ESO QC THRUPUT"] = thr_1m

          for exten in range(17,33,2):
            data = hdu_tell[exten].data
            if data == None:
              hdu_tell[exten].header["HIERARCH ESO QC ZPOINT"] = zpt_2m
              hdu_tell[exten].header["HIERARCH ESO QC THRUPUT"] = thr_2m

          for exten in range(33,49,2):
            data = hdu_tell[exten].data
            if data == None:
              hdu_tell[exten].header["HIERARCH ESO QC ZPOINT"] = zpt_3m
              hdu_tell[exten].header["HIERARCH ESO QC THRUPUT"] = thr_3m


#  print zpt_1m
#  print zpt_2m
#  print zpt_3m
  
#ACT ACCORDING TO strategy
  output_datasetname=in_sof.datasetName
  files = list()

  # STRATEGY = 0: USE THE INPUT TELLURIC.FITS FILE
  #broadcast the input TELLURIC file 
  if strategy == 0 and IsTelluricPresent == 1:
     files.append(reflex.FitsFile(name_telluric,'TELLURIC',None,output_purpose))
     newsof = reflex.SetOfFiles(output_datasetname,files)
     outputs.out_sof = newsof
  # file requested but not present, and not even the response is present, then broadcasting empy sof
  if strategy == 0 and IsTelluricPresent == 0 and IsResponsePresent == 0:
     newsof = reflex.SetOfFiles(output_datasetname,files)
     outputs.out_sof = newsof

  # file requested but not present, but Response is present. Broadcasting response.
  if strategy == 0 and IsTelluricPresent == 0 and IsResponsePresent == 1:
     files.append(reflex.FitsFile(name_response,'TELLURIC',None,output_purpose))
     newsof = reflex.SetOfFiles(output_datasetname,files)
     outputs.out_sof = newsof


  # STRATEGY = 1: USE THE INPUT TELLURIC_only.FITS FILE

  
  # file requested but not present, broadcasting empy sof
  if strategy ==1 and IsTelluricOnlyPresent == 0:
     newsof = reflex.SetOfFiles(output_datasetname,files)
     outputs.out_sof = newsof
     
  # file is present broadcast it (independently of the presence of a TELLURIC file)
  if strategy == 1 and IsTelluricOnlyPresent == 1:
     files.append(reflex.FitsFile(name_telluric_only,'TELLURIC',None,output_purpose))
     newsof = reflex.SetOfFiles(output_datasetname,files)
     outputs.out_sof = newsof
     

  # STRATEGY = 2: USE THE RESPONSE TELLURIC_only.FITS FILE
  # file requested but not present, broadcasting empy sof
  if strategy ==2 and IsResponsePresent == 0:
     newsof = reflex.SetOfFiles(output_datasetname,files)
     outputs.out_sof = newsof
  # file is present and telluric not (=header unchanged).
  if strategy ==2 and IsResponsePresent == 1 and IsTelluricPresent == 0:
     files.append(reflex.FitsFile(name_response,'TELLURIC',None,output_purpose))
     newsof = reflex.SetOfFiles(output_datasetname,files)
     outputs.out_sof = newsof
  # file is present and also telluric  (=update zpt in the header)
  if strategy ==2 and IsResponsePresent == 1 and IsTelluricPresent == 1:
     hdu_response = fits.open(name_response)
     new_name_r = output_dir + '/RESPONSE_ZPT_UPDATED.fits'

     
     for exten in range(1,17,1):
        try:
           hdu_response[exten].header["HIERARCH ESO QC ZPOINT"] = hdu_tell[exten].header["HIERARCH ESO QC ZPOINT"]
           hdu_response[exten].header["HIERARCH ESO QC THRUPUT"] = hdu_tell[exten].header["HIERARCH ESO QC THRUPUT"]
        except KeyError:
           hdu_response[exten].header["HIERARCH ESO QC ZPOINT"] = zpt_1m
           hdu_response[exten].header["HIERARCH ESO QC THRUPUT"] =thr_1m

     for exten in range(17,33,1):
        try:
           hdu_response[exten].header["HIERARCH ESO QC ZPOINT"] = hdu_tell[exten].header["HIERARCH ESO QC ZPOINT"]
           hdu_response[exten].header["HIERARCH ESO QC THRUPUT"] = hdu_tell[exten].header["HIERARCH ESO QC THRUPUT"]
        except KeyError:
           hdu_response[exten].header["HIERARCH ESO QC ZPOINT"] = zpt_2m
           hdu_response[exten].header["HIERARCH ESO QC THRUPUT"] =thr_2m

     for exten in range(33,49,1):
        try:
           hdu_response[exten].header["HIERARCH ESO QC ZPOINT"] = hdu_tell[exten].header["HIERARCH ESO QC ZPOINT"]
           hdu_response[exten].header["HIERARCH ESO QC THRUPUT"] = hdu_tell[exten].header["HIERARCH ESO QC THRUPUT"]
        except KeyError:
           hdu_response[exten].header["HIERARCH ESO QC ZPOINT"] = zpt_3m
           hdu_response[exten].header["HIERARCH ESO QC THRUPUT"] =thr_3m


     hdu_response.writeto(new_name_r,clobber=True)
     files.append(reflex.FitsFile(new_name_r,'TELLURIC',None,output_purpose))
     newsof = reflex.SetOfFiles(output_datasetname,files)
     outputs.out_sof = newsof

  # STRATEGY = 3: COMBINE RESPONSE AND TELLURIC_ONLY
  # IF at least one file is missing, I broadcast an empy sof
  if strategy == 3 and (IsTelluricOnlyPresent == 0 or IsResponsePresent == 0): 
     newsof = reflex.SetOfFiles(output_datasetname,files)
     outputs.out_sof = newsof


  # If the needed files are present  
  if strategy == 3 and IsTelluricOnlyPresent == 1 and IsResponsePresent == 1: 
     hdu_response = fits.open(name_response)
     hdu_tell_only = fits.open(name_telluric_only)
     new_name_rt = output_dir + '/TELLURIC_AND_RESPONSE.fits'
     hdu_tell_only[0].header["HIERARCH ESO PRO CATG"]="'TELLURIC'"
  # I use the ZPT information from TELLURIC_ONLY. It is supposed to be computed either on science data or STD data, which both have
  # the correct ZPT information. If this info is missing, I copy the one in the RESPONSE file.

     for exten in range(1,39,1):
        try:
           wk1= hdu_tell_only[exten].header["HIERARCH ESO QC ZPOINT"]
           wk2= hdu_tell_only[exten].header["HIERARCH ESO QC THRUPUT"]
        except KeyError:
           val1=hdu_response[exten].header["HIERARCH ESO QC ZPOINT"]
           val2=hdu_response[exten].header["HIERARCH ESO QC THRUPUT"]
           hdu_tell_only[exten].header.append(('HIERARCH ESO QC ZPOINT',val1,'median IFU zeropoint'),end=True)#["HIERARCH ESO QC ZPOINT"] = hdu_response[exten].header["HIERARCH ESO QC ZPOINT"]
           hdu_tell_only[exten].header.append(('HIERARCH ESO QC THRUPUT',val2,'median IFU throughput'),end=True)#["HIERARCH ESO QC THRUPUT"] = hdu_response[exten].header["HIERARCH ESO QC THRUPUT"]
     # 
     # Now I can combine response and telluric_only. The file will contain as many valid extensions as TELLURIC_ONLY
     for exten in range(1,49,2):
           data = hdu_tell_only[exten].data
           noise = hdu_tell_only[exten+1].data
           response_ =  hdu_response[exten].data
           error_response_ = hdu_response[exten+1].data
      
           if data != None:
              wdata = getw(hdu_tell_only[exten].header)
              wresponse = getw(hdu_response[exten].header)
              #I need to interpolate the response and error_response to the wavelength range and step of data and noise.
              response = np.interp(wdata,wresponse,response_)
              hdu_tell_only[exten].data = data * response
           if noise != None:
              wnoise = wget(hdu_res[exten+1].header)
              #I need to interpolate the response and error_response to the wavelength range and step of data and noise.
              error_response =  np.interp(wnoise,wresponse,error_response_)
              hdu_tell_only[exten+1].data = (noise**2 +  error_response**2 )**0.5
             
          
     #output_purpose=file.purposes       
     #output_category.append('TELLURIC')
     hdu_tell_only.writeto(new_name_rt,clobber=True)
     files.append(reflex.FitsFile(new_name_rt,'TELLURIC',None,output_purpose))
     newsof = reflex.SetOfFiles(output_datasetname,files)
     outputs.out_sof = newsof


  # print outputs
  parser.write_outputs()
  sys.exit()
