# import the needed modules
from __future__ import absolute_import
from __future__ import print_function
try:
  import reflex
  import pipeline_product
  import_sucess = True

#NOTE for developers: 
# -If you want to modify the current script to cope
#  with different parameters, this is the function to modify:
#  setInteractiveParameters()
# -If you want to modify the current script to read different data from
#  the input FITS, this is the function to modify:
#  readFitsData()                  (from class DataPlotterManager) 
# -If you want to modify the current script to modify the plots (using the same
#  data),  this is the function to modify:
#  plotProductsGraphics()          (from class DataPlotterManager)
# -If you want to modify the text that appears in the "Help" button,
#  this is the function to modify:
#  setWindowHelp()
# -If you want to modify the title of the window, modify this function:
#  setWindowTitle()


  #This class deals with the specific details of data reading and final plotting.
  class DataPlotterManager:
    # This function will read all the columns, images and whatever is needed
    # from the products. The variables , self.plot_x, self.plot_y, etc...
    # are used later in function plotProductsGraphics().
    # Add/delete these variables as you need (only that plotProductsGraphics()
    # has to use the same names).
    # You can also create some additional variables (like statistics) after
    # reading the files.
    # If you use a control variable (self.xxx_found), you can modify 
    # later on the layout of the plotting window based on the presence of 
    # given input files. 
    # sof contains all the set of frames
    def readFitsData(self, fitsFiles):
      #Control variable to check if the interesting files where at the input
      self.red_sci_point_2d_found = False
      self.red_sci_point_blue_found = False
      self.fluxcal_sci_point_blue_found = False
      self.red_sci_point_red_found = False
      self.fluxcal_sci_point_red_found = False
      otype = "POINT_"
      #Read all the products
      frames = dict()
      for frame in fitsFiles:
        if frame == '': continue
        try:
          pp = pipeline_product.PipelineProduct(frame)
          pro_rec1_id = pp.readKeyword('ESO PRO REC1 ID')
        except KeyError:
          pro_rec1_id = '<missing>'
        if pro_rec1_id == 'uves_obs_scired':
          print('Keeping {} with PRO.REC1.ID={}'.format(frame.category, pro_rec1_id))
        else:
          print('Tossing {} with PRO.REC1.ID={}'.format(frame.category, pro_rec1_id))
          continue
        category = frame.category
        frames[category] = frame

      #print(frames.keys())
      if [x for x in frames.keys() if x.startswith("RED_SCI_POINT")]:
        otype = "POINT_"
      elif [x for x in frames.keys() if x.startswith("RED_SCI_SLICER")]:
        otype = "SLICER_"
      elif [x for x in frames.keys() if x.startswith("RED_SCI_EXTND")]:
        otype = "EXTND_"
      elif [x for x in frames.keys() if x.startswith("RED_2D_SCIENCE")]: 
        self.red_sci_point_2d_found = True
      elif [x for x in frames.keys() if x.startswith("RED_2D_SCI_POINT")]: 
        self.red_sci_point_2d_found = True
      elif [x for x in frames.keys() if x.startswith("RED_2D_SCI_EXTND")]: 
        self.red_sci_point_2d_found = True

      self.otype=otype
      print(otype)

      key1 = "RED_SCI_"+otype+"BLUE"
      key2 = "ERRORBAR_SCI_"+otype+"BLUE"
      key3 = "ORDER_EXTRACT_QC_BLUE"
      if key1 in frames and \
         key2 in frames and \
         key3 in frames :
        self.red_sci_point_blue_found = True
        self.spec_raw = uves_plot_common.UvesBlueSpectrum(frames["RED_SCI_"+otype+"BLUE"],
                                         frames["ERRORBAR_SCI_"+otype+"BLUE"])
        self.order_table = uves_plot_common.PlotableBlueOrderTable(frames["ORDER_EXTRACT_QC_BLUE"])

        key4 = "FLUXCAL_SCI_"+otype+"BLUE"
        key5 = "FLUXCAL_ERRORBAR_SCI_"+otype+"BLUE"
        if key4 in frames and \
           key5 in frames : 
          self.fluxcal_sci_point_blue_found = True
          self.spec_fluxcal  = uves_plot_common.UvesBlueSpectrum(frames["FLUXCAL_SCI_"+otype+"BLUE"],
                                                frames["FLUXCAL_ERRORBAR_SCI_"+otype+"BLUE"])
        

      key1 = "RED_SCI_"+otype+"REDL"
      key2 = "RED_SCI_"+otype+"REDU"
      key3 = "ERRORBAR_SCI_"+otype+"REDL"
      key4 = "ERRORBAR_SCI_"+otype+"REDU"
      key5 = "ORDER_EXTRACT_QC_REDL"
      key6 = "ORDER_EXTRACT_QC_REDU"
      if key1 in frames and \
         key2 in frames and \
         key3 in frames and \
         key4 in frames and \
         key5 in frames and \
         key6 in frames : 
        self.red_sci_point_red_found = True
        self.spec_raw  = uves_plot_common.UvesRedSpectrum(frames["RED_SCI_"+otype+"REDL"],
                                         frames["RED_SCI_"+otype+"REDU"],
                                         frames["ERRORBAR_SCI_"+otype+"REDL"],
                                         frames["ERRORBAR_SCI_"+otype+"REDU"])
        self.order_table = uves_plot_common.PlotableRedOrderTable(frames["ORDER_EXTRACT_QC_REDL"],
                                                 frames["ORDER_EXTRACT_QC_REDU"])


        key1 = "FLUXCAL_SCI_"+otype+"REDL"
        key2 = "FLUXCAL_SCI_"+otype+"REDU"
        key3 = "FLUXCAL_ERRORBAR_SCI_"+otype+"REDL"
        key4 = "FLUXCAL_ERRORBAR_SCI_"+otype+"REDU"
        if key1 in frames and \
           key2 in frames and \
           key3 in frames and \
           key4 in frames :
          self.fluxcal_sci_point_red_found = True
          self.spec_fluxcal = uves_plot_common.UvesRedSpectrum(frames["FLUXCAL_SCI_"+otype+"REDL"],
                                              frames["FLUXCAL_SCI_"+otype+"REDU"],
                                              frames["FLUXCAL_ERRORBAR_SCI_"+otype+"REDL"],
                                              frames["FLUXCAL_ERRORBAR_SCI_"+otype+"REDU"])

    # This function creates all the subplots. It is responsible for the plotting 
    # layouts. 
    # There can different layouts, depending on the availability of data
    # Note that subplot(I,J,K) means the Kth plot in a IxJ grid 
    # Note also that the last one is actually a box with text, no graphs.
    def addSubplots(self, figure):
      if self.red_sci_point_blue_found == True or \
         self.red_sci_point_red_found == True :
        if self.fluxcal_sci_point_blue_found == True or \
           self.fluxcal_sci_point_red_found == True :
          self.subplot_spec     = figure.add_subplot(3,1,1)
          self.subplot_spec_cal = figure.add_subplot(3,1,2)
          self.subplot_sn       = figure.add_subplot(6,2,9)
          self.subplot_fwhm     = figure.add_subplot(6,2,10)
          self.subplot_ripple   = figure.add_subplot(6,2,11)
          self.subplot_linepos  = figure.add_subplot(6,2,12)
        else :
          self.subplot_spec     = figure.add_subplot(3,1,1)
          self.subplot_sn       = figure.add_subplot(3,2,3)
          self.subplot_fwhm     = figure.add_subplot(3,2,4)
          self.subplot_ripple   = figure.add_subplot(3,2,5)
          self.subplot_linepos  = figure.add_subplot(3,2,6)
      else : 
        self.subtext_nodata     = figure.add_subplot(1,1,1)
          
    # This is the function that makes the plots.
    # Add new plots or delete them using the given scheme.
    # The data has been already stored in self.plot_x, self.plot_xdif, etc ...
    # It is mandatory to add a tooltip variable to each subplot.
    # One might be tempted to merge addSubplots() and plotProductsGraphics().
    # There is a reason not to do it: addSubplots() is called only once at
    # startup, while plotProductsGraphics() is called always there is a resize.
    def plotProductsGraphics(self):
      if self.red_sci_point_blue_found == True or self.red_sci_point_red_found == True :
        #First subpanel: a spectrum
        if      self.red_sci_point_blue_found == True :
          title_spec   = 'Extracted and Merged Spectrum. No. Orders: %g   Slit Length (pix): %#.3g'% (self.spec_raw.qc_ex_nord, self.spec_raw.qc_ex_ysize)
        elif self.red_sci_point_red_found == True :
          title_spec   = 'Extracted and Merged Spectrum. RedLo: No. Orders: %g   Slit Length (pix): %#.3g. RedHi: No. Orders: %g   Slit Length (pix): %#.3g'% (self.spec_raw.qc_ex_nord_low, self.spec_raw.qc_ex_ysize_low, self.spec_raw.qc_ex_nord_high, self.spec_raw.qc_ex_ysize_high) 
        tooltip_spec = """Plot of the extracted and merged spectrum of the science object (blue line) as total flux (ADU) versus wavelength (Ang). 
The +-1 sigma uncertainties are plotted as the light blue region encompassing the object spectrum (and bounded by the light grey spectra). 
Note that this spectrum is not flux calibrated."""
        self.spec_raw.plot(self.subplot_spec, title_spec, tooltip_spec)

        if self.fluxcal_sci_point_blue_found == True or self.fluxcal_sci_point_red_found == True :
          if   self.red_sci_point_blue_found == True :
            title_spec_fluxcal = 'Flux Calibrated Spectrum. No. Orders: %g   Slit Length (pix): %#.3g'% (self.spec_raw.qc_ex_nord, self.spec_raw.qc_ex_ysize) 
            tooltip_spec_fluxcal="""Plot of the extracted, merged and flux calibrated spectrum of the science object (blue line) as flux ("""+self.spec_fluxcal.spectrum.bunit+""") versus wavelength (Ang). 
The +-1 sigma uncertainties are plotted as the light blue region encompassing the object spectrum (and bounded by the light grey spectra)."""
          elif self.red_sci_point_red_found == True :
            title_spec_fluxcal = 'Flux Calibrated Spectrum. RedLo: No. Orders: %g   Slit Length (pix): %#.3g. RedHi: No. Orders: %g   Slit Length (pix): %#.3g'% (self.spec_raw.qc_ex_nord_low, self.spec_raw.qc_ex_ysize_low, self.spec_raw.qc_ex_nord_high, self.spec_raw.qc_ex_ysize_high) 
            tooltip_spec_fluxcal="""Plot of the extracted, merged and flux calibrated spectrum of the science object (blue line) as flux ("""+self.spec_fluxcal.spec_low.bunit+""") versus wavelength (Ang). 
The +-1 sigma uncertainties are plotted as the light blue region encompassing the object spectrum (and bounded by the light grey spectra)."""
          self.spec_fluxcal.plot(self.subplot_spec_cal, 
                                 title_spec_fluxcal, tooltip_spec_fluxcal)

        #Second subpanel: a plot
        title_sn   = 'Min S/N: %#.3g   Max S/N: %#.3g'% (self.order_table.minSN(), self.order_table.maxSN()) 
        tooltip_sn = """Plot of the measured S/N of the extracted and merged spectrum as a function of the order."""
        self.order_table.plotSN(self.subplot_sn, title_sn, tooltip_sn)

        #Third subpanel: a plot
        title_fwhm = 'Min FWHM: %#.3g   Max FWHM: %#.3g'% (self.order_table.minFWHM(), self.order_table.maxFWHM()) 
        tooltip_fwhm = """Plot of the spatial FWHM (pix) of the science object spectrum trace as a function of the order."""
        self.order_table.plotFWHM(self.subplot_fwhm, title_fwhm, tooltip_fwhm)

        #Fourth subpanel: a plot
        title_ripple   = 'Min Ripple: %#.3g   Max Ripple: %#.3g'% (self.order_table.minRipple(), self.order_table.maxRipple()) 
        tooltip_ripple = """Plot of the amplitude of systematic ripples in the science spectrum as a function of the order. 
Negative values indicate where the value of this statistic could not be calculated."""
        self.order_table.plotRipple(self.subplot_ripple, title_ripple, tooltip_ripple)

        #Fifth subpanel: a plot
        title_linepos   = 'Min Pos: %#.3g   Max Pos: %#.3g'% (self.order_table.minLinepos(), self.order_table.maxLinepos())
        tooltip_linepos = """Plot of the position along the slit (pix) of the science object spectrum trace as a function of the order."""
        self.order_table.plotLinepos(self.subplot_linepos, title_linepos, tooltip_linepos)
        if      self.red_sci_point_blue_found == True :
          self.subplot_linepos.set_ylim(0., self.spec_raw.qc_ex_ysize)
        elif self.red_sci_point_red_found == True :
          self.subplot_linepos.set_ylim(0., self.spec_raw.qc_ex_ysize_low)

      elif self.red_sci_point_2d_found == True :
        #2D extraction info
        self.subtext_nodata.set_axis_off()
        self.text_nodata = """Science object spectra extracted with method "2D",
Results are not displayed but can be found in the 
reflex_end_products directory.
The user may still change reduction parameters and 
re-run the recipe or continue."""
        self.subtext_nodata.text(0.1, 0.6, self.text_nodata, color='#11557c', fontsize=18,
                                 ha='left', va='center', alpha=1.0)
        self.subtext_nodata.tooltip='Science object spectra extracted with method 2D'

      else :
        #Data not found info
        self.subtext_nodata.set_axis_off()
        self.text_nodata = """Science object spectra not found in the products:
For Blue data:  PRO.CATG=RED_SCI_%sBLUE
For Red data:  PRO.CATG=RED_SCI_%sREDL, RED_SCI_%sREDU"""%(self.otype,self.otype,self.otype)
        self.subtext_nodata.text(0.1, 0.6, self.text_nodata, color='#11557c', fontsize=18,
                                 ha='left', va='center', alpha=1.0)
        self.subtext_nodata.tooltip='Science object spectra not found in the products'

  
    # This function specifies which are the parameters that should be presented
    # in the window to be edited.
    # Note that the parameter has to be also in the in_sop port (otherwise it 
    # won't appear in the window) 
    # The descriptions are used to show a tooltip. They should match one to one
    # with the parameter list 
    # Note also that parameters have to be prefixed by the 'recipe name:'
    def setInteractiveParameters(self):
      paramList = list()
      paramList.append(reflex.RecipeParameter(recipe='uves_obs_scired',displayName='reduce.extract.method',group='extract',description='Extraction method. <average | linear | 2d | optimal>'))
      paramList.append(reflex.RecipeParameter(recipe='uves_obs_scired',displayName='reduce.extract.kappa',group='extract',description='In optimal extraction mode, this is the threshold for bad (i.e. hot/cold) pixel rejection. If a pixel deviates more than kappa*sigma (where sigma is the uncertainty of the pixel flux) from the inferred spatial profile, its weight is set to zero. Range: [-1,100]. If this parameter is negative, no rejection is performed'))
      paramList.append(reflex.RecipeParameter(recipe='uves_obs_scired',displayName='reduce.extract.profile',group='extract',description='In optimal extraction mode, the kind of profile to use. gauss gives a Gaussian profile, moffat gives a Moffat profile with beta=4 and a possible linear sky contribution. virtual uses a virtual resampling algorithm (i.e. measures and uses the actual object profile). constant assumes a constant spatial profile and allows optimal extraction of wavelength calibration frames. auto will automatically select the best method based on the estimated S/N of the object. For low S/N, moffat or gauss are recommended (for robustness). For high S/N, virtual is recommended (for accuracy). In the case of virtual resampling, a precise determination of the order positions is required; therefore the order-definition is repeated using the (assumed non-low S/N) science frame. <constant | gauss | moffat | virtual | auto>'))
      paramList.append(reflex.RecipeParameter(recipe='uves_obs_scired',displayName='reduce.extract.skymethod',group='extract',description='In optimal extraction mode, the sky subtraction method to use. median estimates the sky as the median of pixels along the slit (ignoring pixels close to the object), whereas optimal does a chi square minimization along the slit to obtain the best combined object and sky levels. The optimal method gives the most accurate sky determination but is also a bit slower than the median method. <median | optimal>'))
      paramList.append(reflex.RecipeParameter(recipe='uves_obs_scired',displayName='reduce.extract.oversample',group='extract',description='The oversampling factor used for the virtual resampling algorithm. If negative, the value 5 is used for S/N <=200, and the value 10 is used if the estimated S/N is > 200'))
      paramList.append(reflex.RecipeParameter(recipe='uves_obs_scired',displayName='reduce.slitlength',group='other',description='Extraction slit length (in pixels). If negative, the value inferred from the raw frame header is used'))
      paramList.append(reflex.RecipeParameter(recipe='uves_obs_scired',displayName='reduce.objoffset',group='other',description='Offset (in pixels) of extraction slit with respect to center of order. This parameter applies to linear/average/optimal extraction. For linear/average extraction, if the related parameter objslit is negative, the offset is automatically determined by measuring the actual object position'))
      paramList.append(reflex.RecipeParameter(recipe='uves_obs_scired',displayName='reduce.objslit',group='other',description='Object window size (in pixels). This must be less than the total slit length. If negative, the default value (half of full slit length) is used. The upper and lower sky windows are defined as the part of the full slit (if any) outside the object window. The center of the object window is determined by the offset parameter.  This parameter does not apply to optimal extraction.'))
      paramList.append(reflex.RecipeParameter(recipe='uves_obs_scired',displayName='reduce.rebin.wavestep',group='rebinning',description='The bin size used for BLUE/REDL data (in w.l.u.) in wavelength space. If negative, a step size of 2/3 * ( average pixel size ) is used'))
      paramList.append(reflex.RecipeParameter(recipe='uves_obs_scired',displayName='reduce.rebin.wavestep_redu',group='rebinning',description='The bin size used for REDU data (in w.l.u.) in wavelength space. If negative, a step size of 2/3 * ( average pixel size ) is used'))

      return paramList

    def setWindowHelp(self):
      help_text = """
In this window, the user should check that the science object extracted spectrum is of a good quality by using the pan and zoom buttons at the top-left of this window. 
Attempt to optimise the S/N of the extracted spectrum as a function of spectral order (the upper plot of the bottom-left plots) by choosing different parameter values and re-running the pipeline recipe."""
      return help_text

    def setWindowTitle(self):
      title = 'Uves Interactive Spectrum Extraction'
      return title

except ImportError:
  import_sucess = 'false'
  print("Error importing modules pyfits, wx, matplotlib, numpy")

#This is the 'main' function
if __name__ == '__main__':

  # import reflex modules
  import reflex_interactive_app
  import sys

  # import UVES reflex modules
  import uves_plot_common

  # Create interactive application
  interactive_app = reflex_interactive_app.PipelineInteractiveApp(enable_init_sop=True)

  #Check if import failed or not
  if import_sucess == 'false' :
    interactive_app.setEnableGUI('false')

  #Open the interactive window if enabled
  if interactive_app.isGUIEnabled() :
    #Get the specific functions for this window
    dataPlotManager = DataPlotterManager()
    interactive_app.setPlotManager(dataPlotManager)
    interactive_app.showGUI()
  else :
    interactive_app.passProductsThrough()

  # print outputs
  interactive_app.print_outputs()

  sys.exit()

