#! /usr/bin/python3

"""\
%prog [<ipolfile>=ipol.dat [<refdir>]]

Interactive parameterisation explorer with optional ref data display

TODO:
 * Need to allow no-ref visualisation
 * Optionally display the envelope of input MC runs
"""

import optparse, os, sys
op = optparse.OptionParser(usage=__doc__)
op.add_option("-v", "--debug", dest="DEBUG", action="store_true", default=False, help="Turn on some debug messages")
op.add_option("-n", dest="NUMWIN", metavar="NUM", type=int, default=4, help="number of windows to generate [default=%default]")
op.add_option("-q", "--quiet", dest="QUIET", action="store_true", default=False, help="Turn off messages")
op.add_option("--wfile", dest="WFILE", default=None, help="Path to a weight file, used to restrict plotting to a subset of histograms (default: %default)")
op.add_option("-t", dest="TFILE", default=None, help="Path to a tune parameter file, used to set initial conditions")
# TODO: Add weight file parsing to decide which histos (and bin subsets) to interpolate
opts, args = op.parse_args()

## Get arguments
IFILE = "ipol.dat"
if len(args) >= 1:
    IFILE = args[0]
REFDIR = None
if len(args) >= 2:
    REFDIR = args[1]


## Load the Professor machinery
import professor2 as prof
from professor2 import misc

if not opts.QUIET:
    print(prof.logo)


## Read in ipol histos
IHISTOS, METADATA = prof.read_ipoldata(IFILE)

## Weight file parsing
if opts.WFILE:
    matchers = prof.read_pointmatchers(opts.WFILE)
    for hn in list(IHISTOS.keys()):
        if not any(m.match_path(hn) for m in list(matchers.keys())):
            del IHISTOS[hn]
    if len(list(IHISTOS.keys()))==0:
        print("Nothing left after weight file parsing, exiting")
        sys.exit(0)

## List of rivet analysis paths in IHISTOS
# TODO: Urgh, avoid this sort of noise-code!
ananames = list(set([[x for x in i.split("/") if len(x)>1][0] for i in list(IHISTOS.keys())]))

## Read reference data histos
import glob
HISTOS = {}
# TODO: consistent directory convention and protect against invalid structures
if REFDIR is not None:
    reffiles = glob.glob(os.path.join(REFDIR, "*"))
    for rf in reffiles:
        HISTOS.update(prof.read_histos(rf))

    ## Find things available in both
    available = []
    for i in list(IHISTOS.keys()):
        if i in list(HISTOS.keys()):
            available.append([i,i])
    ibins = []
    databins = []
    for a in available:
        ibins.extend(  IHISTOS[a[0]].bins)
        databins.extend(HISTOS[a[1]].bins)

else:
    ibins = []
    available = [(i,i) for i in list(IHISTOS.keys())]
    for a in available:
        ibins.extend(  IHISTOS[a[0]].bins)

## Sanity checks

if not ibins:
    print("No bins ..., exiting")
    import sys
    sys.exit(1)


## Import and set up matplotlib (AFTER parsing arguments)
try:
    import matplotlib
    matplotlib.use('WXAgg')
    from matplotlib.figure import Figure
    from matplotlib.backends.backend_wxagg import FigureCanvasWxAgg as FigCanvas
    from matplotlib.backends.backend_wxagg import NavigationToolbar2WxAgg as NavigationToolbar
    import wx, functools
    params = {
        'backend': 'pdf',
        'axes.labelsize' : 16,
        'text.size' : 16,
        'legend.fontsize' : 16,
        'axes.titlesize' : 16,
        'xtick.labelsize' : 16,
        'ytick.labelsize' : 16,
        'text.usetex' : False,
        'figure.dpi' : 50,
        'lines.markersize' : 10,
        'lines.linewidth' : 2,
        'lines.elinewidth' : 3,
        'lines.antialiased' : False,
        #'patches.antialiased' : False,
        'figure.subplot.left' : 0.05,
        'figure.subplot.right' : 0.995,
        'figure.subplot.bottom' : 0.1,
        'figure.subplot.top' : 0.95,
        'figure.subplot.wspace' : 0.15
        }
    for k, v in list(params.items()):
        try:
            matplotlib.rcParams[k] = v
        except:
            pass
except Exception as e:
    print("Problem with getting & configuring matplotlib/WX interface method: %s" % e)
    print("Exiting!")
    sys.exit(1)

observables=sorted([a[0] for a in available]) # When interested in comparison with data
#observables=sorted(IHISTOS.keys())            # all ipols, not neccessary data available


# TODO: add error bar plotting option (with errs from ipol)


def mk_line(obs, ppoint, derivative=False):
    """
    Prepare points for matplotlib step
    """
    X, Y = [], []
    for b in IHISTOS[obs].bins:
        X.append(b.xmin)
        X.append(b.xmax)
        if derivative:
            y = b.der(ppoint)
        else:
            y = b.val(ppoint)
        Y.append(y)
        Y.append(y)
    return X, Y

def mk_data(obs, prefix=""):
    X, Y, dY = [], [], []
    try:
        for b in HISTOS[prefix+obs].bins:
            X.append(b.xmid)
            Y.append(b.val)
            dY.append(b.err)
    except Exception as e:
        if opts.DEBUG:
            print("No data available for %s"%obs)
    return X, Y, dY

def mk_pull(obs, ppoint, prefix="/REF"):
    """
    Prepare points for matplotlib step
    """
    print("called mk_pull")
    X, Y = [], []
    for num, b in enumerate(IHISTOS[obs].bins):
        X.append(b.xmin)
        X.append(b.xmax)
        p=prof.misc.pull(HISTOS[prefix+obs].bins[num], b, ppoint)
        Y.append(p)
        Y.append(p)
    return X, Y


class ParamFrame(wx.Frame):
    """
    Simple frame with param sliders that update the main frame's plots.
    """
    def __init__(self, parent):
        no_sys_menu = wx.CAPTION
        wx.Frame.__init__(self, None,  title='Prof-I: Parameter control')#, style=no_sys_menu)
        self.params = METADATA["ParamNames"].split()
        self.panel = wx.Panel(self)
        self.parent = parent
        self.sliders = {}

        ## Weight file parsing
        self.tunePars={}
        if opts.TFILE:
            try:
                with open(opts.TFILE) as f:
                    for lin in f:
                        l=lin.split()
                        #  print(l)
                        self.tunePars[l[0]]=float(l[1])
                        #  print(self.tunePars.keys())
                        #  print(self.tunePars[l[0]])
            except:
                print(f"problem with file {opts.TFILE}")
                sys.exit(0)
        ## Parameter sliders and buttons
        controlgrid = wx.GridBagSizer(hgap=5, vgap=2)

        font = wx.Font(10, wx.MODERN, wx.NORMAL, wx.BOLD)
        ## Build the sliders and related objects
        for i, param in enumerate(self.params):
            ### Add param names
            t=wx.StaticText(self.panel, -1, param)
            t.SetFont(font)
            controlgrid.Add(t, pos=(i,0))
            ## Sliders
            #http://www.wxpython.org/docs/api/wx.lib.agw.floatspin-module.html
            from wx.lib.agw.floatspin import FloatSpin, EVT_FLOATSPIN
            pmin = float(METADATA["MinParamVals"].split()[i])
            pmax = float(METADATA["MaxParamVals"].split()[i])

            from wx.lib.agw.floatspin import FloatSpin, EVT_FLOATSPIN
            #  print(param)
            #  print(self.tunePars[param])
            if opts.TFILE is not None:
                try:
                    print(f"{param} tuned to {self.tunePars[param]}")
                except:
                    print(f"Problem with {opts.TFILE}")
                    sys.exit(2)
            else:
                self.tunePars[param]=pmin + 0.5*(pmax-pmin)
            slider = FloatSpin(self.panel, -1, style=0, value=self.tunePars[param],
            #  slider = FloatSpin(self.panel, -1, style=0, value=pmin + 0.5*(pmax-pmin),
                               min_val=pmin, max_val=pmax, increment=abs(pmax-pmin)/50., size=(400,40))
            controlgrid.Add(slider, pos=(i,2))
            self.sliders[param] = slider
        self.Bind(EVT_FLOATSPIN, self.parent.sliderUpdate)
        self.panel.SetSizer(controlgrid)
        controlgrid.Fit(self)


class ProfIFrame(wx.Frame):

    def __init__(self):
        title = 'prof-I: Professor Interactive'
        wx.Frame.__init__(self, None, -1, title)
        # The Frame with the parameter thingies
        self.child = ParamFrame(self)
        self.child.Show()
        # This frame
        self.histos = {}
        self.axes = []
        self.obschoice = []
        self.logxchkbx = []
        self.logychkbx = []
        self.derchkbx = []
        self.pullchkbx = []
        self.resetBtn = []
        self.refhistos = {}
        self.tunehisto = {}
        self.logy = {}
        self.logx = {}
        self.pull = {}
        self.derivative = {}
        self.limits = {}
        self.drawobs = { }
        print(f"Observable number {len(observables)}")
        print(f"Number of windows {opts.NUMWIN}")
        self.createMainPanel()
        #  self.updateObs(None, "left")
        #  self.updateObs(None, "right")


    def createMainPanel(self):
        """ Creates the main panel with all the controls on it:
             * mpl canvas
             * mpl navigation toolbar
             * Control panel for interaction
        """
        self.dpi=50
        self.panel = wx.Panel(self)

        ## Panel for observable diplay
        self.fig = Figure((10.0, 4.0), dpi=self.dpi, facecolor="white")
        self.canvas = FigCanvas(self.panel, -1, self.fig)
        NmaxCol=int(opts.NUMWIN/3)
        for i in range(opts.NUMWIN):
            #  if colidx == NmaxCol:
                #  colidx=0
                #  rowidx+=1
            #  print(f"{rowidx},{colidx},{i}")
            #  self.axes.append(self.fig.add_subplot(rowidx+1,colidx+1,i+1))
            self.axes.append(self.fig.add_subplot(3,NmaxCol,i+1))
            #  colidx+=1

        ## Layout with box sizers
        ## matplotlib canvas and matplotlib toolbar
        self.vbox = wx.BoxSizer(wx.VERTICAL)
        self.vbox.Add(self.canvas, 1, wx.LEFT | wx.TOP | wx.GROW | wx.EXPAND)
        self.hbox = wx.GridBagSizer(hgap=5, vgap=2)
        
        iidx=0
        rowidx=0
        colidx=0
        for i in range(opts.NUMWIN):
            if colidx == NmaxCol:
                colidx=0
                rowidx+=1
            self.drawobs[f"Obs{i}"]=sorted(observables)[i if len(observables) > i else len(observables)-1]
            #  print(f"{rowidx},{colidx}")
            #  self.hbox.Add(wx.StaticText(self.panel, -1, f"Obs {i}:"), pos=(0,iidx))
            self.hbox.Add(wx.StaticText(self.panel, -1, f"Obs {i}:"), pos=(rowidx,6*colidx))
            self.obschoice.append(wx.Choice(self.panel, -1, (85, 18), choices=sorted(observables)))
            self.obschoice[i].SetSelection(sorted(observables).index(self.drawobs[f"Obs{i}"]))
            #  self.hbox.Add(self.obschoice[i], pos=(0,iidx+1))
            self.hbox.Add(self.obschoice[i], pos=(rowidx,6*colidx+1))
            self.logychkbx.append(wx.CheckBox(self.panel, -1, "logy"))
            self.logxchkbx.append(wx.CheckBox(self.panel, -1, "logx"))
            #  self.hbox.Add(self.logychkbx[i], pos=(0,iidx+2))
            #  self.hbox.Add(self.logxchkbx[i], pos=(0,iidx+3))
            self.hbox.Add(self.logychkbx[i], pos=(rowidx,6*colidx+2))
            self.hbox.Add(self.logxchkbx[i], pos=(rowidx,6*colidx+3))
            self.Bind(wx.EVT_CHOICE, functools.partial(self.updateObs, which=f"Obs{i}"), self.obschoice[i])
            self.Bind(wx.EVT_CHECKBOX, self.setLogy, self.logychkbx[i])
            self.Bind(wx.EVT_CHECKBOX, self.setLogx, self.logxchkbx[i])
            self.pullchkbx.append(wx.CheckBox(self.panel, -1, "pull"))
            #  self.hbox.Add(self.pullchkbx[i], pos=(0,iidx+4))
            self.hbox.Add(self.pullchkbx[i], pos=(rowidx,6*colidx+4))
            self.derchkbx.append(wx.CheckBox(self.panel, -1, "derivative"))
            #  self.hbox.Add(self.derchkbx[i], pos=(0,iidx+5))
            self.hbox.Add(self.derchkbx[i], pos=(rowidx,6*colidx+5))
            self.Bind(wx.EVT_CHOICE, functools.partial(self.updateObs, which=f"Obs{i}"), self.obschoice[i])
            self.Bind(wx.EVT_CHECKBOX, self.setPull, self.pullchkbx[i])
            self.Bind(wx.EVT_CHECKBOX, self.setDerivative, self.derchkbx[i])
            self.refhistos[f"Obs{i}"]=None
            self.tunehisto[f"Obs{i}"]=None
            self.logx[f"Obs{i}"]=False
            self.logy[f"Obs{i}"]=False
            self.pull[f"Obs{i}"]=False
            self.derivative[f"Obs{i}"]=False
            self.updateObs(None, f"Obs{i}")
            iidx+=6
            colidx+=1

        self.hbox1 = wx.BoxSizer(wx.HORIZONTAL)

        ## vbox for x/y limits
        self.limitgrid = wx.GridBagSizer(hgap=5, vgap=2)
        # Buttons to reset
        limittypes = ("XMin", "XMax", "YMin", "YMax")
        idxY=0
        for i in range(opts.NUMWIN):
            self.limitgrid.Add(wx.StaticText(self.panel, -1, f"Lim {i}:"), pos=(0,idxY))
            self.resetBtn.append(wx.Button(self.panel, label="Reset"))
            self.Bind(wx.EVT_BUTTON, functools.partial(self.resetLimits, which=f"Obs{i}"), self.resetBtn[i])
            self.limitgrid.Add(self.resetBtn[i], pos=(0,idxY+1))
            self.limits[f"Obs{i}"] = {}
            for num, ind in enumerate(limittypes):
                self.limitgrid.Add(wx.StaticText(self.panel, -1, ind), pos=(num+1, idxY))
                textctrl = wx.TextCtrl(self.panel, -1, "None")
                self.limitgrid.Add(textctrl, pos=(num+1,idxY+1))
                self.limits[f"Obs{i}"][ind] = textctrl
            for textctrl in list(self.limits[f"Obs{i}"].values()):
                self.Bind(wx.EVT_TEXT, self.sliderUpdate, textctrl)
            idxY+=3

        #  self.resetBtn2 = wx.Button(self.panel, label="Reset")
        #  self.Bind(wx.EVT_BUTTON, functools.partial(self.resetLimits, which="right"), self.resetBtn2)
        #  self.limitgrid.Add(self.resetBtn2, pos=(0,4))
        # Input fields

        #  self.limits["left"] = {}
        #  for num, i in enumerate(limittypes):
            #  self.limitgrid.Add(wx.StaticText(self.panel, -1, i), pos=(num+1, 0))
            #  textctrl = wx.TextCtrl(self.panel, -1, "None")
            #  self.limitgrid.Add(textctrl, pos=(num+1,1))
            #  self.limits["left"][i] = textctrl
        #  self.limitgrid.Add(wx.StaticText(self.panel, -1, "Limits 2:"), pos=(0,3))

        #  self.limits["right"] = {}
        #  for num, i in enumerate(limittypes):
            #  self.limitgrid.Add(wx.StaticText(self.panel, -1, i), pos=(num+1, 3))
            #  textctrl = wx.TextCtrl(self.panel, -1, "None", style=wx.TE_PROCESS_ENTER)
            #  self.limitgrid.Add(textctrl, pos=(num+1,4))
            #  self.limits["right"][i] = textctrl
        self.hbox.Fit(self)

        self.vbox.Add(self.hbox)
        #self.vbox.AddSpacer((10,10))
        self.vbox.AddSpacer(10)
        self.hbox1.Add(self.limitgrid)


        ## Add  to the main vertical box
        self.vbox.Add(self.hbox1)

        self.panel.SetSizer(self.vbox)
        self.vbox.Fit(self)

        ## Maximize the window
        self.Maximize()


    def setLogy(self, event):
        """ See if the checkboxes are checked and set logy-scales accordingly """
        arr=[]
        for i in range(len(self.drawobs.keys())):
            arr.append([self.logychkbx[i], f"Obs{i}", self.axes[i]])
        for i in arr:
        #  for i in [(self.logy1chkbx, "left", self.axes), (self.logy2chkbx, "right", self.axes2)]:
            if i[0].IsChecked():
                #self.logy[i[1]]=True
                i[2].set_yscale('log')
            else:
                #self.logy[i[1]]=False
                i[2].set_yscale('linear')
        self.sliderUpdate(None)

    def setLogx(self, event):
        """ See if the checkboxes are checked and set logx-scales accordingly """
        arr=[]
        for i in range(len(self.drawobs.keys())):
            arr.append([self.logxchkbx[i], f"Obs{i}", self.axes[i]])
        for i in arr:
        #  for i in [(self.logx1chkbx, "left", self.axes), (self.logx2chkbx, "right", self.axes2)]:
            if i[0].IsChecked():
                #self.logx[i[1]]=True
                i[2].set_xscale('log')
            else:
                #self.logx[i[1]]=False
                i[2].set_xscale('linear')
        self.sliderUpdate(None)

    def setPull(self, event):
        """ See if the checkboxes are checked and switch to pull mode """
        arr=[]
        for i in range(len(self.drawobs.keys())):
            arr.append([self.pullchkbx[i], f"Obs{i}", self.axes[i]])
        for i in arr:
        #  for i in [(self.pull1chkbx, "left", self.axes), (self.pull2chkbx, "right", self.axes2)]:
            if i[0].IsChecked():
                self.pull[i[1]]=True
            else:
                self.pull[i[1]]=False
        self.sliderUpdate(None)

    def setDerivative(self, event):
        """ See if the checkboxes are checked and switch to pull mode """
        arr=[]
        for i in range(len(self.drawobs.keys())):
            arr.append([self.derchkbx[i], f"Obs{i}", self.axes[i]])
        for i in arr:
        #  for i in [(self.der1chkbx, "left", self.axes), (self.der2chkbx, "right", self.axes2)]:
            if i[0].IsChecked():
                self.derivative[i[1]]=True
                self.updateObs(None,i[1])
            else:
                self.derivative[i[1]]=False
                self.updateObs(None,i[1])
        self.sliderUpdate(None)

    def sliderUpdate(self, event, redraw=True):
        #  print(f"called with {redraw}")
        ppoint = self.getCurrentParamValues()
        self.plotHistos(ppoint, redraw=redraw)
        ## Set plot limits
        for i in range(len(self.limits.keys())):
            self.setLimits(self.axes[i], self.getLimits(self.limits[f"Obs{i}"]))

        #  self.setLimits(self.axes2, self.getLimits(self.limits["right"]))
        self.canvas.draw()

    def getCurrentParamValues(self):
        return [self.child.sliders[k].GetValue() for k in self.child.params]

    def getTuneParamValues(self):
        return [self.child.tunePars[k] for k in self.child.params]

    def resetLimits(self, event, which):
        for i in "XMin XMax YMin YMax".split():
            self.limits[which][i].SetValue("None")
        self.updateObs(None, which)

    def getLimits(self, limdict):
        xmin = self.convertInput(limdict["XMin"].GetValue())
        xmax = self.convertInput(limdict["XMax"].GetValue())
        ymin = self.convertInput(limdict["YMin"].GetValue())
        ymax = self.convertInput(limdict["YMax"].GetValue())
        return xmin, xmax, ymin, ymax

    def setLimits(self, sub, thelimits):
        xmin, xmax, ymin, ymax = thelimits
        sub.set_xlim(xmin, xmax)
        sub.set_ylim(ymin, ymax)


    def convertInput(self, inp):
        try:
            return float(inp)
        except:
            return None

    def updateObs(self, event=None, which=None):
        """ Update the observable assigned the L or R window.
        """
        if which is not None:
            idx=int(which[3:])
            for ii in range(len(self.obschoice)):
                self.axes[ii].clear()
                if not self.pull[f"Obs{ii}"]:
                    self.plotDataHisto(self.axes[ii], f"Obs{ii}")
            #  print(which)
            #  print(idx)
            #  print("obsch")
            #  print(self.obschoice)
            #  print(self.obschoice)
            #  print(self.obschoice[idx].GetSelection())
            self.drawobs[which] = observables[self.obschoice[idx].GetSelection()]
            obs= self.drawobs[which]
            self.axes[idx].set_title(obs)
            #  if opts.NUMWIN == len(self.obschoice):
                #  self.sliderUpdate(None, redraw=True)
            #  else:
            self.sliderUpdate(None, redraw=False)
            #  self.sliderUpdate(None, redraw=True)
            return 
        for i in range(opts.NUMWIN):
            self.axes[i].clear()
            self.drawobs[f"Obs{i}"] = observables[self.obschoice[i].GetSelection()]
            obs= self.drawobs[f"Obs{i}"]
            self.axes[i].set_title(obs)
            #  if not self.pull[f"Obs{i}"]:
                #  self.plotDataHisto(self.axes[i], f"Obs{i}")

        #  self.axes.set_title(obs1)
        #  self.axes2.set_title(obs2)

        #  if not self.pull['left']:
            #  self.plotDataHisto(self.axes, 'left')
        #  if not self.pull['right']:
            #  self.plotDataHisto(self.axes2, 'right')



        #  self.axes.clear()
        #  self.axes2.clear()
        #  if which in ("left", None):
            #  self.drawobs["left"] = observables[self.obschoice1.GetSelection()]
        #  if which in ("right", None):
            #  self.drawobs["right"] = observables[self.obschoice2.GetSelection()]

        # Set titles and axis labels, using the rivet plot file parser if possible
        #  obs1 = self.drawobs["left"]
        #  obs2 = self.drawobs["right"]
        #  self.axes.set_title(obs1)
        #  self.axes2.set_title(obs2)

        #  if not self.pull['left']:
            #  self.plotDataHisto(self.axes, 'left')
        #  if not self.pull['right']:
            #  self.plotDataHisto(self.axes2, 'right')

        self.sliderUpdate(None, redraw=False)
        #  self.sliderUpdate(None, redraw=True)

    def plotDataHisto(self, sub, which):
        x,y,ye = mk_data(self.drawobs[which])
        ppoint=self.getTuneParamValues()
        i=int(which[3:])
        X, Y = mk_line(self.drawobs[which], ppoint, self.derivative[which])#, self.logy[i[0]])
        self.tunehisto[which] = self.axes[i].step(X, Y, color="b", where="post")
        if len(x)>0:
            self.refhistos[which] = sub.errorbar(x,y, yerr=ye, ls=" ", marker="o", color="k", label="Data")


    def plotHistos(self, ppoint, redraw=False):
        """ This plots the observable prediction calculated from the
            parameterisation.
        """
        #  print(f"called with {redraw}")
        arr=[]
        for i in range(len(self.drawobs.keys())):
            arr.append([f"Obs{i}", self.axes[i], self.drawobs[f"Obs{i}"]])
        for i in arr:

            # Get the coordinates for the ipol histos step draw
            if self.pull[i[0]]:
                X, Y = mk_pull(i[2], ppoint)#, self.logy[i[0]])
            else:
                X, Y = mk_line(i[2], ppoint, self.derivative[i[0]])#, self.logy[i[0]])
            if redraw is False:
                self.histos[i[0]] = i[1].step(X, Y, color="r", where="post")
                #  print(self.histos[i[0]])
            else:
                self.histos[i[0]][0].set_xdata(X)
                self.histos[i[0]][0].set_ydata(Y)
                #  print(self.histos[i[0]][0])
                ## Workaround for changes in matplotlib 0.99
                try:
                    i[1].redraw_in_frame()
                except:
                    pass

## Run GUI application
app = wx.App()
ProfIFrame().Show()
app.MainLoop()
