"""
    Dialog panel to explore the P(r) inversion results for a range
    of D_max value. User picks a number of points and a range of
    distances, then can toggle between inversion outputs and see
    their distribution as a function of D_max.
    
    This software was developed by the University of Tennessee as part of the
    Distributed Data Analysis of Neutron Scattering Experiments (DANSE)
    project funded by the US National Science Foundation. 

    See the license text in license.txt

    copyright 2009, University of Tennessee
"""
import wx
import numpy
import math
import logging
import sys

# Avoid Matplotlib complaining about the lack of legend on the plot 
import warnings
warnings.simplefilter("ignore")

# Import plotting classes
from danse.common.plottools.PlotPanel import PlotPanel
from danse.common.plottools import Theory1D as Model1D
from danse.common.plottools.plottables import Graph

from pr_widgets import PrTextCtrl

# Default number of points on the output plot
DEFAULT_NPTS = 10
# Default output parameter to plot
DEFAULT_OUTPUT = 'Chi2/dof'

class OutputPlot(PlotPanel):
    """
        Plot panel used to show the selected results as a function
        of D_max
    """
    def __init__(self, d_min, d_max, parent, id = -1, color = None,\
        dpi = None, style = wx.NO_FULL_REPAINT_ON_RESIZE, **kwargs):
        """
            Initialization. The parameters added to PlotPanel are:
            
            @param d_min: Minimum value of D_max to explore
            @param d_max: Maximum value of D_max to explore
        """
        PlotPanel.__init__(self, parent, id = id, style = style, **kwargs)

        self.parent = parent
        self.min = d_min
        self.max = d_max
        self.npts = DEFAULT_NPTS

        step = (self.max-self.min)/(self.npts-1)
        self.x = numpy.arange(self.min, self.max+step*0.01, step)
        dx = numpy.zeros(len(self.x))
        y = numpy.ones(len(self.x))
        dy = numpy.zeros(len(self.x))
 
        # Plot area
        self.plot = Model1D(self.x,y=y,dy=dy)
        self.plot.name = DEFAULT_OUTPUT

        # Graph        
        self.graph = Graph()
        self.graph.xaxis("\\rm{D_{max}}", 'A')
        self.graph.yaxis("\\rm{%s}" % DEFAULT_OUTPUT,"")
        self.graph.add(self.plot)
        self.graph.render(self)

    def onContextMenu(self, event):
        """
            Default context menu for the plot panel
            
            #TODO: Would be nice to add printing and log/linear scales.
            The current verison of plottools no longer plays well with 
            plots outside of guiframe. Guiframe team needs to fix this. 
        """
        # Slicer plot popup menu
        id = wx.NewId()
        slicerpop = wx.Menu()
        slicerpop.Append(id,'&Save image', 'Save image as PNG')
        wx.EVT_MENU(self, id, self.onSaveImage)
        
        id = wx.NewId()
        slicerpop.AppendSeparator()
        slicerpop.Append(id, '&Reset Graph')
        wx.EVT_MENU(self, id, self.onResetGraph)
       
        pos = event.GetPosition()
        pos = self.ScreenToClient(pos)
        self.PopupMenu(slicerpop, pos)

        
class Results:
    """
        Class to hold the inversion output parameters
        as a function of D_max
    """
    def __init__(self):
        """
            Initialization. Create empty arrays
            and dictionary of labels.
        """
        # Array of output for each inversion
        self.chi2 = []
        self.osc = []
        self.pos = []
        self.pos_err = []
        self.rg = []
        self.iq0 = []
        self.bck = []
        self.d_max = []
        
        # Dictionary of outputs
        self.outputs = {}
        self.outputs['Chi2/dof'] = ["\chi^2/dof", "a.u.", self.chi2]    
        self.outputs['Oscillation parameter'] = ["Osc", "a.u.", self.osc]    
        self.outputs['Positive fraction'] = ["P^+", "a.u.", self.pos]    
        self.outputs['1-sigma positive fraction'] = ["P^+_{1\ \sigma}", "a.u.", self.pos_err]    
        self.outputs['Rg'] = ["R_g", "A", self.rg]    
        self.outputs['I(q=0)'] = ["I(q=0)", "1/A", self.iq0]    
        self.outputs['Background'] = ["Bck", "1/A", self.bck] 
        
class ExploreDialog(wx.Dialog):
    """
        The explorer dialog box. This dialog is meant to be
        invoked by the InversionControl class.
    """
    
    def __init__(self, pr_state, nfunc, *args, **kwds):
        """
            Initialization. The parameters added to Dialog are:
            
            @param pr_state: sans.pr.invertor.Invertor object
            @param nfunc: Number of terms in the expansion
        """
        kwds["style"] = wx.RESIZE_BORDER|wx.DEFAULT_DIALOG_STYLE
        wx.Dialog.__init__(self, *args, **kwds)
        
        # Initialize Results object
        self.results = Results()
        
        self.pr_state  = pr_state
        self._default_min = 0.9*self.pr_state.d_max
        self._default_max = 1.1*self.pr_state.d_max
        self.nfunc     = nfunc
        
        # Control for number of points
        self.npts_ctl = PrTextCtrl(self, -1, style=wx.TE_PROCESS_ENTER, size=(60,20))
        # Control for the minimum value of D_max
        self.dmin_ctl = PrTextCtrl(self, -1, style=wx.TE_PROCESS_ENTER, size=(60,20))
        # Control for the maximum value of D_max
        self.dmax_ctl = PrTextCtrl(self, -1, style=wx.TE_PROCESS_ENTER, size=(60,20))

        # Output selection box for the y axis
        self.output_box = None

        # Create the plot object
        self.plotpanel = OutputPlot(self._default_min, self._default_max, self, -1, style=wx.RAISED_BORDER)
        
        # Create the layout of the dialog
        self.__do_layout()
        self.Fit()
        
        # Calculate exploration results
        self._recalc()
        # Graph the default output curve
        self._plot_output()
        
    class Event:
        """
            Class that holds the content of the form
        """
        ## Number of points to be plotted
        npts   = 0
        ## Minimum value of D_max
        dmin   = 0
        ## Maximum value of D_max
        dmax   = 0
                
    def _get_values(self, event=None):
        """
            Invoked when the user changes a value of the form.
            Check that the values are of the right type.
            
            @return: ExploreDialog.Event object if the content is good, None otherwise
        """
        # Flag to make sure that all values are good
        flag = True
        
        # Empty ExploreDialog.Event content
        content_event = self.Event()
        
        # Read each text control and make sure the type is valid
        # Let the user know if a type is invalid by changing the
        # background color of the control.
        try:
            content_event.npts = int(self.npts_ctl.GetValue())
            self.npts_ctl.SetBackgroundColour(wx.WHITE)
            self.npts_ctl.Refresh()
        except:
            flag = False
            self.npts_ctl.SetBackgroundColour("pink")
            self.npts_ctl.Refresh()
            
        try:
            content_event.dmin = float(self.dmin_ctl.GetValue())
            self.dmin_ctl.SetBackgroundColour(wx.WHITE)
            self.dmin_ctl.Refresh()
        except:
            flag = False
            self.dmin_ctl.SetBackgroundColour("pink")
            self.dmin_ctl.Refresh()
        
        try:
            content_event.dmax = float(self.dmax_ctl.GetValue())
            self.dmax_ctl.SetBackgroundColour(wx.WHITE)
            self.dmax_ctl.Refresh()
        except:
            flag = False
            self.dmax_ctl.SetBackgroundColour("pink")
            self.dmax_ctl.Refresh()
        
        # If the content of the form is valid, return the content,
        # otherwise return None
        if flag:
            if event is not None:
                event.Skip(True)
            return content_event
        else:
            return None
    
    def _plot_output(self, event=None):
        """
            Invoked when a new output type is selected for plotting, 
            or when a new computation is finished.
        """
        # Get the output type selection
        output_type = self.output_box.GetString(self.output_box.GetSelection())
        
        # If the selected output type is part of the results ojbect,
        # display the results. 
        # Note: by design, the output type should always be part of the
        #       results object.
        if self.results.outputs.has_key(output_type): 
            self.plotpanel.plot.x = self.results.d_max
            self.plotpanel.plot.y = self.results.outputs[output_type][2]
            self.plotpanel.plot.name = '_nolegend_'
            self.plotpanel.graph.yaxis("\\rm{%s}" % self.results.outputs[output_type][0], self.results.outputs[output_type][1])
            
            # Redraw
            self.plotpanel.graph.render(self.plotpanel)
            self.plotpanel.subplot.figure.canvas.draw_idle()
        else:
            msg =  "ExploreDialog: the Results object's dictionary does not contain "
            msg += "the [%s] output type. This must be indicative of a change in the " % str(output_type)
            msg += "ExploreDialog code."
            logging.error(msg)
        
    def __do_layout(self):
        """
            Do the layout of the dialog
        """
        # Dialog box properties
        self.SetTitle("D_max Explorer")
        self.SetSize((600, 595))
        
        sizer_main = wx.BoxSizer(wx.VERTICAL)
        sizer_button = wx.BoxSizer(wx.HORIZONTAL)
        sizer_params = wx.GridBagSizer(5,5)

        iy = 0
        ix = 0
        label_npts  = wx.StaticText(self, -1, "Npts")
        sizer_params.Add(label_npts, (iy,ix), (1,1), wx.LEFT|wx.EXPAND|wx.ADJUST_MINSIZE, 15)
        ix += 1
        sizer_params.Add(self.npts_ctl,   (iy,ix), (1,1), wx.EXPAND|wx.ADJUST_MINSIZE, 0)
        self.npts_ctl.SetValue("%g" % DEFAULT_NPTS)
        self.npts_ctl.Bind(wx.EVT_KILL_FOCUS, self._recalc)
        
        ix += 1
        label_dmin   = wx.StaticText(self, -1, "Min Distance [A]")
        sizer_params.Add(label_dmin, (iy,ix), (1,1), wx.LEFT|wx.EXPAND|wx.ADJUST_MINSIZE, 15)
        ix += 1
        sizer_params.Add(self.dmin_ctl,   (iy,ix), (1,1), wx.EXPAND|wx.ADJUST_MINSIZE, 0)
        self.dmin_ctl.SetValue(str(self._default_min))
        self.dmin_ctl.Bind(wx.EVT_KILL_FOCUS, self._recalc)
        
        ix += 1
        label_dmax = wx.StaticText(self, -1, "Max Distance [A]")
        sizer_params.Add(label_dmax, (iy,ix), (1,1), wx.LEFT|wx.EXPAND|wx.ADJUST_MINSIZE, 15)
        ix += 1
        sizer_params.Add(self.dmax_ctl,   (iy,ix), (1,1), wx.EXPAND|wx.ADJUST_MINSIZE, 0)
        self.dmax_ctl.SetValue(str(self._default_max))
        self.dmax_ctl.Bind(wx.EVT_KILL_FOCUS, self._recalc)


        # Ouput selection box
        selection_msg = wx.StaticText(self, -1, "Select a dependent variable:")
        self.output_box = wx.ComboBox(self, -1)
        for item in self.results.outputs.keys():
            self.output_box.Append(item, "")
        self.output_box.SetStringSelection(DEFAULT_OUTPUT)
        
        output_sizer = wx.GridBagSizer(5,5)
        output_sizer.Add(selection_msg, (0,0), (1,1), wx.LEFT|wx.EXPAND|wx.ADJUST_MINSIZE, 10)
        output_sizer.Add(self.output_box, (0,1), (1,2), wx.LEFT|wx.EXPAND|wx.ADJUST_MINSIZE, 10)
        
        wx.EVT_COMBOBOX(self.output_box,-1, self._plot_output) 
        sizer_main.Add(output_sizer, 0, wx.EXPAND | wx.ALL, 10)
        
        sizer_main.Add(self.plotpanel, 0, wx.EXPAND | wx.ALL, 10)
        sizer_main.SetItemMinSize(self.plotpanel, 400,400)
        
        sizer_main.Add(sizer_params, 0, wx.EXPAND|wx.ALL, 10)
        static_line_3 = wx.StaticLine(self, -1)
        sizer_main.Add(static_line_3, 0, wx.EXPAND, 0)
        
        # Bottom area with the close button
        sizer_button.Add((20, 20), 1, wx.EXPAND|wx.ADJUST_MINSIZE, 0)
        button_OK = wx.Button(self, wx.ID_OK, "Close")
        sizer_button.Add(button_OK, 0, wx.LEFT|wx.RIGHT|wx.ADJUST_MINSIZE, 10)
        
        sizer_main.Add(sizer_button, 0, wx.EXPAND|wx.BOTTOM|wx.TOP, 10)
        self.SetAutoLayout(True)
        self.SetSizer(sizer_main)
        self.Layout()
        self.Centre()
        
        # Bind the Enter key to recalculation
        self.Bind(wx.EVT_TEXT_ENTER, self._recalc)

    def _recalc(self, event=None):
        """
            Invoked when the user changed a value on the form.
            Process the form and compute the output to be plottted.
        """
        # Get the content of the form
        content = self._get_values()
        # If the content of the form is invalid, return and do nothing
        if content is None:
            return
        
        # Results object to store the computation outputs.
        results = Results()
        
        # Loop over d_max values
        for i in range(content.npts):    
            d = content.dmin + i * (content.dmax - content.dmin)/(content.npts-1.0)
            self.pr_state.d_max = d
            try:
                out, cov = self.pr_state.invert(self.nfunc)   
            
                # Store results
                iq0 = self.pr_state.iq0(out)
                rg = self.pr_state.rg(out)
                pos = self.pr_state.get_positive(out)
                pos_err = self.pr_state.get_pos_err(out, cov)
                osc = self.pr_state.oscillations(out)
            
                results.d_max.append(self.pr_state.d_max)
                results.bck.append(self.pr_state.background)
                results.chi2.append(self.pr_state.chi2)
                results.iq0.append(iq0)
                results.rg.append(rg)
                results.pos.append(pos)
                results.pos_err.append(pos_err)
                results.osc.append(osc)           
            except:
                # This inversion failed, skip this D_max value
                logging.error("ExploreDialog: inversion failed for D_max=%s\n%s" % (str(d), sys.exc_value))
            
        self.results = results            
         
        # Plot the selected output
        self._plot_output()   
            
##### testing code ############################################################
class MyApp(wx.App):
    """
        Test application used to invoke the ExploreDialog for testing
    """
    def OnInit(self):
        from inversion_state import Reader
        from sans.pr.invertor import Invertor
        wx.InitAllImageHandlers()
        
        def call_back(state, datainfo=None):
            """
                Dummy call-back method used by the P(r) 
                file reader.
            """
            print state
            
        # Load test data
        # This would normally be loaded by the application 
        # of the P(r) plug-in.
        r = Reader(call_back)
        data = r.read("test_pr_data.prv")
        pr_state = data.meta_data['prstate']
        
        # Create Invertor object for the data
        # This would normally be assembled by the P(r) plug-in
        pr = Invertor()
        pr.d_max = pr_state.d_max
        pr.alpha = pr_state.alpha
        pr.q_min = pr_state.qmin
        pr.q_max = pr_state.qmax
        pr.x = data.x
        pr.y = data.y
        pr.err = data.y*0.1
        pr.has_bck = False
        pr.slit_height = pr_state.height
        pr.slit_width = pr_state.width
        
        # Create the dialog
        dialog = ExploreDialog(pr, 10, None, -1, "")
        self.SetTopWindow(dialog)
        dialog.ShowModal()
        dialog.Destroy()

        return 1

if __name__ == "__main__":
    app = MyApp(0)
    app.MainLoop()