# import the needed modules
try:
  import reflex
  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
        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. (2d/optimal not supported by uves_cal_wavecal, weighted supported only by uves_cal_wavecal). <average | linear | 2d | weighted | 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.rebin.wavestep',group='rebinning',description='The bin size (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()

  #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()

