import reflex 
import sys 
import os

try:
    import pyfits as fits
except:
    from astropy.io import fits 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__':
  parser = reflex.ReflexIOParser()
  parser.add_input("-i", "--star_spc_in", dest="star_spc_in")
  parser.add_input("-j", "--response_in", dest="response_in")
  parser.add_output("-o", "--out_sof", dest="out_sof")

  inputs = parser.get_inputs()

  inputs_spc  = inputs.star_spc_in
  inputs_res  = inputs.response_in
  outputs = parser.get_outputs()

  files_res=inputs_res.files
  files_spc=inputs_spc.files

  #get the output directory
  pattern = '--products-dir'
  for arg in sys.argv:
     if arg.split("=")[0] == pattern:
         output_dir = arg.split("=")[1]

  new_name=list()
  output_category=list()


 
  #START PROCEDURE
  
  #read response if present
  IsResponsePresent = 0
  for file in files_res:
      hdu_res = fits.open(file.name)
      IsResponsePresent = 1
  
  
  #  if respone is present, do some stuff and create the output sof
  if IsResponsePresent == 1:
    c = 1
    for file in files_spc:
      
      hdu_spc = fits.open(file.name)
      newname = output_dir + '/STD_RESP_CORR_'+str(c)+'.fits'
      new_name.append(newname)
      
      for exten in range(1,49,2):
           data = hdu_spc[exten].data
           noise = hdu_spc[exten+1].data
           response_ =  hdu_res[exten].data
           error_response_ = hdu_res[exten+1].data
      
           if data is not None:
              
              wdata = getw(hdu_spc[exten].header)
              wresponse = getw(hdu_res[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_spc[exten].data = data / response
              
           if noise is not None:
              wnoise = getw(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_spc[exten+1].data = ((noise/response)**2 + (data/response**2 * error_response)**2 )**0.5
              
      
      output_purpose=file.purposes       
      output_category.append(file.category)
      hdu_spc.writeto(newname,clobber=True)
      c = c+1

    #WRITE OUTPUT SOF
    files = list()
    output_datasetname=inputs_spc.datasetName
    for i in range(len(new_name)):
       files.append(reflex.FitsFile(new_name[i],output_category[i],None,output_purpose))
    newsof = reflex.SetOfFiles(output_datasetname,files)
    outputs.out_sof = newsof

  # If I did not have a response do nothing; output sof is equal to the input STD  sof 
  else:
    outputs.out_sof = inputs.star_spc_in

  # print outputs
  parser.write_outputs()

  sys.exit()
