from __future__ import with_statement
from __future__ import absolute_import
from __future__ import print_function
import sys


try:
    import numpy
    import reflex
    from pipeline_product import PipelineProduct
    import pipeline_display
    import reflex_plot_widgets
    import matplotlib.gridspec as gridspec
    import_success = True
    import pdb  # for debugging

except ImportError:
    import_success = False
    print("Error importing modules pyfits, wx, matplotlib, numpy")

# Median absolute deviation function; used to scale the images
def MAD(x):
    x=numpy.array(x)
    return numpy.median(numpy.abs(x-numpy.median(x)))


def paragraph(text, width=None):
    """ wrap text string into paragraph
       text:  text to format, removes leading space and newlines
       width: if not None, wraps text, not recommended for tooltips as
              they are wrapped by wxWidgets by default
    """
    import textwrap
    if width is None:
        return textwrap.dedent(text).replace('\n', ' ').strip()
    else:
        return textwrap.fill(textwrap.dedent(text), width=width)


class DataPlotterManager(object):
    """
    This class must be added to the PipelineInteractiveApp with setPlotManager
    It must have following member functions which will be called by the app:
     - setInteractiveParameters(self)
     - readFitsData(self, fitsFiles):
     - addSubplots(self, figure):
     - plotProductsGraphics(self, figure, canvas)
    Following members are optional:
     - setWindowHelp(self)
     - setWindowTitle(self)
    """

    # static members
    recipe_name = "hawki_detector_noise"
    mst_cat = "MASTER_READGAIN"
    ref_cat = "REFERENCE_READGAIN"

    def setWindowTitle(self):
        return self.recipe_name+"_interactive"

    def setInteractiveParameters(self):
        """
        This function specifies which are the parameters that should be presented
        in the window to be edited.  Note that the parameter has to also be 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.
        """
        return [

            reflex.RecipeParameter(recipe=self.recipe_name, displayName="thresh",
                                   group="hawki_detector_noise", 
                                   description="Rejection threshold in sigma above background. [5.0]"),
            ]

    def readFitsData(self, fitsFiles):
        """
        This function should be used to read and organize the raw fits files
        produced by the recipes.
        It receives as input a list of reflex.FitsFiles
        """
        # organize the files into a dictionary, here we assume we only have 
        # one file per category if there are more, one must use a
        # dictionary of lists
        self.frames = dict()
        for f in fitsFiles:
            self.frames[f.category] = PipelineProduct(f)

        # we only have two states, we have data or we don't
        # define the plotting functions we want to use for each

        if self.mst_cat in self.frames:
            self.mst_tab = self.frames[self.mst_cat]
            self.mst_found = True

            # Read the reference table

            if self.ref_cat in self.frames:
                self.ref_tab = self.frames[self.ref_cat]
            else:
                self.ref_found = False

            # table is a list of FITS record arrays, one for each extension
            # access data by field name: table['COLNAME']
            # see help at https://pythonhosted.org/pyfits/users_guide/users_table.html
            
            table = self.mst_tab.all_hdu[1].data
            # check if data in MASTER_READGAIN is valid
            if ( (False in numpy.isfinite(table['READNOISE'])) or
                 (False in numpy.isfinite(table['GAIN'])) or
                 (False in numpy.isfinite(table['COVAR']))  or 
               ( (True in (table['READNOISE'] <= 0)) or 
                 (True in (table['GAIN'] <= 0)) or 
                 (True in (table['COVAR'] <= 0)) ) ):

                # At least one data element is not good, so set the plotting functions to Baddata ones
                self._add_subplots = self._add_nodata_subplots
                self._plot = self._baddata_plot

            else:
                self._add_subplots = self._add_subplots
                self._plot = self._data_plot
                
        else:
            # No master data, so set the plotting functions to NODATA ones
            self._add_subplots = self._add_nodata_subplots
            self._plot = self._nodata_plot

    def addSubplots(self, figure):
        self._add_subplots(figure)

    def _add_subplots(self, figure):

        # Make 1x2 grid of plots
        self.img_plot = []
        gs = gridspec.GridSpec(1, 2)
        self.img_plot.append(figure.add_subplot(gs[0,0]))
        self.img_plot.append(figure.add_subplot(gs[0,1]))

    def plotProductsGraphics(self):
        self._plot()

    def _data_plot(self):
        
        title = " " 
        tooltip = "If a REFERENCE value equals a derived MASTER value, the point will lie on the solid blue line"

        colours = ['red','blue','green','purple']
        markers = ["o","s","^","D"]

        mst_table = self.mst_tab.all_hdu[1].data
        ref_table = self.ref_tab.all_hdu[1].data
        
        for i in range(2):
            self.img_plot[i].grid(True)
            self.img_plot[i].set_title(title, fontsize=12, fontweight='semibold')
            self.img_plot[i].tooltip = tooltip
            #self.img_plot[i].set_aspect('equal')
            self.img_plot[i].set_aspect('auto')

        # Pull out correct numbers in table, be careful to match reference chip to master chip
        chip_mst = mst_table['EXTNAME']
        chip_ref = ref_table['EXTNAME']
        x_left=[]
        y_left=[]
        x_right=[]
        y_right=[]
        chip_label = []
        for i in range(len(chip_ref)):
            for j in range(len(chip_mst)):
                if ( chip_mst[j] == chip_ref[i]):
                    x_left.append(mst_table['READNOISE'][j])
                    y_left.append(ref_table['READNOISE'][i])
                    x_right.append((mst_table['GAIN'][j]) / (mst_table['COVAR'][j]))
                    y_right.append((ref_table['GAIN'][i]) / (ref_table['COVAR'][i]))
                    chip_label.append(chip_mst[j])

        self.img_plot[0].set_xlabel("READNOISE [ADU]")
        self.img_plot[0].set_ylabel("REFERENCE READNOISE [ADU]")
        
        for i_row in range(len(mst_table['EXTNAME'])):
            self.img_plot[0].scatter(x_left[i_row], y_left[i_row], 60, color = colours[i_row],marker = markers[i_row], label = chip_label[i_row])

        self.img_plot[0].plot([min(x_left),max(x_left)],[min(x_left),max(x_left)], label="REF == MASTER")
        self.img_plot[0].legend(loc = 0, scatterpoints = 1)

        cur_ylim = self.img_plot[0].get_ylim()
        self.img_plot[0].set_ylim([cur_ylim[0], cur_ylim[1]+0.2*(cur_ylim[1]-cur_ylim[0])])

        self.img_plot[1].set_xlabel("GAIN [e-/ADU]")
        self.img_plot[1].set_ylabel("REFERENCE GAIN [e-/ADU]")
        
        for i_row in range(len(mst_table['EXTNAME'])):
            self.img_plot[1].scatter(x_right[i_row], y_right[i_row], 60, color = colours[i_row],marker = markers[i_row], label = chip_label[i_row])
        self.img_plot[1].plot([min(x_right),max(x_right)],[min(x_right),max(x_right)], label="REF == MASTER")
        self.img_plot[1].legend(loc = 0, scatterpoints = 1)

        cur_ylim = self.img_plot[1].get_ylim()
        self.img_plot[1].set_ylim([cur_ylim[0], cur_ylim[1]+0.2*(cur_ylim[1]-cur_ylim[0])])

    def _add_nodata_subplots(self, figure):
        self.img_plot = figure.add_subplot(1,1,1)

    def _nodata_plot(self):
        # could be moved to reflex library?
        self.img_plot.set_axis_off()
        text_nodata = "Data not found. Input files should contain this" \
                       " type:\n%s" % self.mst_cat
        self.img_plot.text(0.1, 0.6, text_nodata, color='#11557c',
                      fontsize=18, ha='left', va='center', alpha=1.0)
        self.img_plot.tooltip = 'No data found'

    def _baddata_plot(self):
        # could be moved to reflex library?
        self.img_plot.set_axis_off()
        text_baddata = "At least one value in MASTER_READGAIN table\n is out of bounds\n" \
                      "Check the log and input files"
        self.img_plot.text(0.1, 0.6, text_baddata, color='#11557c',
                      fontsize=18, ha='left', va='center', alpha=1.0)
        self.img_plot.tooltip = 'Bad data'

    def setWindowHelp(self):
      help_text = """
This is an interactive window which help asses the quality of the execution of a recipe.
"""
      return help_text

#This is the 'main' function
if __name__ == '__main__':
    from reflex_interactive_app import PipelineInteractiveApp

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

    # get inputs from the command line
    interactive_app.parse_args()

    #Check if import failed or not
    if not import_success:
        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.set_continue_mode()

    #Print outputs. This is parsed by the Reflex python actor to
    #get the results. Do not remove
    interactive_app.print_outputs()
    sys.exit()
