'''
This file contains all the functions that constitute the "frontend of the backend",
i.e. native Python plotting functions. All actual plots are generated by plotting.py --
this is purely about displaying them. manualfit() and geogui() rely on PyQt4, which
is likely to fail, so it's sequestered.

Version: 2019aug06
'''

## Imports and globals...need Qt since matplotlib doesn't support edit boxes, grr!
from optima import OptimaException, dcp, printv, sigfig, makeplots, getplotselections, gridcolors, odict, isnumber, promotetolist, loadobj, sanitizeresults, reanimateplots
from pylab import figure, close, floor, ion, ioff, isinteractive, ceil, array, show, pause
from pylab import subplot, ylabel, transpose, legend, fill_between, xlim, title
from matplotlib.widgets import CheckButtons, Button

global panel, results, origpars, tmppars, parset, fulllabellist, fullkeylist, fullsubkeylist, fulltypelist, fullvallist, plotfig, panelfig, check, checkboxes, updatebutton, clearbutton, defaultsbutton, advancedbutton, closebutton, plusbutton, minusbutton, plotargs, scrwid, scrhei, globaladvanced  # For manualfit GUI
if 1:  panel, results, origpars, tmppars, parset, fulllabellist, fullkeylist, fullsubkeylist, fulltypelist, fullvallist, plotfig, panelfig, check, checkboxes, updatebutton, clearbutton, defaultsbutton, advancedbutton, closebutton, plusbutton, minusbutton, plotargs, scrwid, scrhei, globaladvanced = [None]*25
scrwid, scrhei = 24, 12 # Specify these here...if too large, should shrink anyway


def importpyqt():
    ''' Try to import pyqt, either PyQt4 or PyQt5, but allow it to fail '''
    try:    
        from PyQt4 import QtGui as pyqt
    except:
        try:    from PyQt5 import QtWidgets as pyqt
        except: pyqt = Exception('QtGui could not be imported')
    return  pyqt
    
pyqt = importpyqt()


##############################################################################
### USER-VISIBLE FUNCTIONS
##############################################################################

def plotresults(results, toplot=None, fig=None, figargs=None, **kwargs):
    ''' 
    Does the hard work for updateplots() for pygui()
    Keyword arguments if supplied are passed on to figure().
    
    Usage:
        results = P.runsim('default')
        plotresults(results)
        
    Version: 2016jan25
    '''
    
    if figargs is None: figargs = dict()
    if fig is None: fig = makenewfigure(**figargs)
    
    # Do plotting
    wasinteractive = isinteractive() # You might think you can get rid of this...you can't!
    if wasinteractive: ioff()
    width,height = fig.get_size_inches()
    
    # Actually create plots
    if 'figsize' in kwargs: kwargs.pop('figsize', None)
    plots = makeplots(results, toplot=toplot, die=True, figsize=(width, height), fig=fig, **kwargs)
    naxes = len(plots[0].axes) # If done interactively, they're all in the first plot
    
    # Calculate the rows and columns
    aspectratio = 1.5 # The target aspect ratio for plots, for choosing how many rows and columns to use
    nrows = 1
    ncols = 1
    while nrows*ncols < naxes:
        if width/ncols/aspectratio > height/nrows: ncols += 1 # Height is more squashed: add a column
        else:                                      nrows += 1 # Width is more squashed: add a row
    
    # Adjust margins
    fig.subplots_adjust(left=0.07, bottom=0.05, right=0.85, top=0.95, wspace=0.9, hspace=0.7) # NB, 1.0 seems meaningless for wspace and hspace...
    for a,ax in enumerate(plots[-1].axes):
        ax.change_geometry(nrows, ncols, a+1)

    # Handle interactivity like a boss
    if wasinteractive: ion()
    show()
    return None


def pygui(tmpresults, toplot=None, advanced=False, verbose=2, figargs=None, **kwargs):
    '''
    PYGUI
    
    Make a Python GUI for plotting results. Opens up a control window and a plotting window,
    and when "Update" is clicked, will clear the contents of the plotting window and replot.
    
    Usage:
        pygui(results, [toplot])
    
    where results is the output of e.g. runsim() and toplot is an optional list of form e.g.
        toplot = ['prev-tot', 'inci-pop']
    
    (see epiformatslist in plotting.py)
    
    Warning: the plots won't resize automatically if the figure is resized, but if you click
    "Update", then they will.    
    
    Version: 1.3 (2017feb07)
    '''
    
    global check, checkboxes, updatebutton, clearbutton, defaultsbutton, advancedbutton, closebutton, plusbutton, minusbutton, panelfig, results, plotargs, globaladvanced
    plotargs = kwargs # Reset global to match function input
    results = sanitizeresults(tmpresults)
    globaladvanced = advanced
    
    ## Define options for selection
    plotselections = getplotselections(results, advanced=globaladvanced)
    checkboxes = plotselections['keys']
    checkboxnames = plotselections['names']
    isselected = []
    toplot = promotetolist(toplot) # Ensure it's a list
    if not toplot or toplot[0] is None or toplot[0]=='default': 
        if len(toplot): toplot.pop(0) # Remove the first element
        defaultboxes = [checkboxes[i] for i,tf in enumerate(plotselections['defaults']) if tf] # Back-convert defaults from true/false list to list of keys
        toplot.extend(defaultboxes)
    if len(toplot):
        tmptoplot = dcp(toplot) # Make a copy to compare arguments
        for key in checkboxes:
            if key in toplot:
                isselected.append(True)
                tmptoplot.remove(key)
            else:
                isselected.append(False)
        if len(tmptoplot)!=0:
            errormsg = 'Not all keys were recognized; mismatched ones were:\n'
            errormsg += '%s\n' % tmptoplot
            errormsg += 'Available keys are:\n'
            errormsg += '%s' % checkboxes
            if not globaladvanced: errormsg += '\nSet advanced=True for more options'
            printv(errormsg, 1, verbose=verbose)
    
    ## Set up control panel
    if advanced: 
        figwidth = 14
        advwid = 0.6 # Adjust button width
    else:
        figwidth = 7
        advwid = 1.0
    figheight = 12
    panelfig = figure(num='Optima control panel', figsize=(figwidth,figheight), facecolor=(0.95, 0.95, 0.95)) # Open control panel
    xinit = 0.10*advwid
    if advanced: cbapos = [xinit*advwid,  0.07, 0.9, 1.8] # cba="check box axes position": extra tall, for moving later
    else:        cbapos = [xinit,         0.07, 0.8, 0.9]
    ypos = 0.02 # y-position of buttons
    bwid = 0.14*advwid # x-width of buttons
    bhei = 0.03 # y-height of buttons
    sep  = 0.165 # Separation between buttons
    pmwid = 0.03*advwid # Width of plus/minus buttons
    checkboxaxes = panelfig.add_axes(cbapos) # Create checkbox locations
    updateaxes   = panelfig.add_axes([xinit+0*sep, ypos, bwid, bhei]) # Create update button location
    clearaxes    = panelfig.add_axes([xinit+1*sep, ypos, bwid, bhei]) # Create clear button location
    defaultsaxes = panelfig.add_axes([xinit+2*sep, ypos, bwid, bhei]) # Create defaults button location
    advancedaxes = panelfig.add_axes([xinit+3*sep, ypos, bwid, bhei]) # Create defaults button location
    closeaxes    = panelfig.add_axes([xinit+4*sep, ypos, bwid, bhei]) # Create close button location
    plusaxes     = panelfig.add_axes([xinit+5*sep, ypos+0.015, pmwid, 0.02]) # Create plus button location
    minusaxes    = panelfig.add_axes([xinit+5*sep, ypos-0.005, pmwid, 0.02]) # Create plus button location
    check = CheckButtons(checkboxaxes, checkboxnames, isselected) # Actually create checkboxes
    
    # Reformat the checkboxes
    stastr = ' - stacked'
    perstr = ' - population'
    nboxes = len(check.rectangles)
    for b in range(nboxes):
        label = check.labels[b]
        labeltext = label.get_text()
        labelpos = label.get_position()
        label.set_position((labelpos[0]*0.3,labelpos[1])) # Not sure why by default the check boxes are so far away
        if labeltext.endswith(perstr):    label.set_text('Per population') # Clear label
        elif labeltext.endswith(stastr):  label.set_text('Stacked') # Clear label
        else:                             label.set_weight('bold')
    
    #  If advanced, split into two columns -- messy since Matplotlib sucks! :(
    if advanced: 
        for b in range(nboxes):
            percol = floor(nboxes/2.0) # Number of boxes per column
            col = floor(b/percol) # Which column to plto in
            
            labelpos = list(check.labels[b].get_position()) # Get label positions and convert tuple -> list
            rectpos = list(check.rectangles[b].get_xy()) # Likewise for rectangles
            line0pos = check.lines[b][0].get_data() # There are two lines, and they have data
            line1pos = check.lines[b][1].get_data()
            
            yoffset = 0.5 # Specify amount to move everything in column 0 down by
            xoffset = 0.45 # Specify amount to move everything on column 1 over by
            if col==0: # Left column: shift everything down
                labelpos[1] -= yoffset
                rectpos[1] -= yoffset
                for i in range(2): # Start and end points
                    line0pos[1][i] -= yoffset
                    line1pos[1][i] -= yoffset
            else: # Right column: shift everything over
                labelpos[0] += xoffset
                rectpos[0] += xoffset
                for i in range(2): # Start and end points
                    line0pos[0][i] += xoffset
                    line1pos[0][i] += xoffset
            check.labels[b].set_position(labelpos) # Actually set positions
            check.rectangles[b].set_xy(rectpos)
            check.lines[b][0].set_data(line0pos)
            check.lines[b][1].set_data(line1pos)
    
    if advanced: advlabel = 'Normal'
    else:        advlabel = 'Advanced'
    blue  = (0.4,0.7,1.0) # Also green = (0.2,0.7,0.1), red   = (1.0,0.5,0.1)
    white = (1.0,1.0,1.0)
    black = (0.4,0.4,0.4)
    darker = 0.7
    updatebutton   = Button(updateaxes,   'Update',   color=blue, hovercolor=tuple(array(blue)*darker)) 
    clearbutton    = Button(clearaxes,    'Clear',    color=blue,  hovercolor=tuple(array(blue)*darker))
    defaultsbutton = Button(defaultsaxes, 'Defaults', color=blue,  hovercolor=tuple(array(blue)*darker))
    advancedbutton = Button(advancedaxes,  advlabel,  color=blue,  hovercolor=tuple(array(blue)*darker))
    closebutton    = Button(closeaxes,    'Close',    color=blue,   hovercolor=tuple(array(blue)*darker))
    plusbutton     = Button(plusaxes,      '+',       color=white, hovercolor=tuple(array(white)*darker))
    minusbutton    = Button(minusaxes,     '-',       color=black, hovercolor=tuple(array(black)*darker))
    updatebutton.on_clicked(updateplots) # Update figure if button is clicked
    clearbutton.on_clicked(clearselections) # Clear all checkboxes
    defaultsbutton.on_clicked(defaultselections) # Return to default selections
    advancedbutton.on_clicked(advancedselections) # Return to default selections
    closebutton.on_clicked(closegui) # Close figures
    plusbutton.on_clicked(zoomin) # Zoom in on plots
    minusbutton.on_clicked(zoomout) # Zoom in on plots
    updateplots(None) # Plot initially -- ACTUALLY GENERATES THE PLOTS
    return None




def manualfit(project=None, parsubset=None, name=-1, ind=0, maxrows=25, verbose=2, advanced=False, figargs=None, **kwargs):
    ''' 
    Create a GUI for doing manual fitting via the backend. Opens up three windows: 
    results, results selection, and edit boxes.
    
    parsubset can be a list of parameters the user can fit, e.g.
    parsubset=['initprev','force']
    
    maxrows is the number of rows (i.e. parameters) to display in each column.
    
    Note: to get advanced parameters and plots, set advanced=True.
    
    Version: 1.2 (2017feb10)
    '''
    
    ## Random housekeeping
    global panel, results, origpars, tmppars, parset, fulllabellist, fullkeylist, fullsubkeylist, fulltypelist, fullvallist, globaladvanced
    globaladvanced = advanced
    if figargs is None: figargs = dict()
    fig = figure(**figargs); close(fig) # Open and close figure...dumb, no? Otherwise get "QWidget: Must construct a QApplication before a QPaintDevice"
    ion() # We really need this here!
    nsigfigs = 4
    
    boxes = []
    texts = []
    
    ## Get the list of parameters that can be fitted
    parset = dcp(project.parsets[name])
    tmppars = parset.pars
    origpars = dcp(tmppars)
    
    mflists = parset.manualfitlists(parsubset=parsubset, advanced=globaladvanced)
    fullkeylist    = mflists['keys']
    fullsubkeylist = mflists['subkeys']
    fulltypelist   = mflists['types']
    fullvallist    = mflists['values']
    fulllabellist  = mflists['labels']
    
    nfull = len(fulllabellist) # The total number of boxes needed
    results = project.runsim(name)
    pygui(results, **kwargs)
    
    
    
    def closewindows():
        ''' Close all three open windows '''
        closegui()
        panel.close()
    
    
    ## Define update step
    def manualupdate():
        ''' Update GUI with new results '''
        global results, tmppars, fullkeylist, fullsubkeylist, fulltypelist, fullvallist
        
        # Update parameter values from GUI values
        for b,box in enumerate(boxes): 
            fullvallist[b] = eval(str(box.text())) 
        
        # Create lists for update
        mflists = odict()
        mflists['keys'] = fullkeylist
        mflists['subkeys'] = fullsubkeylist
        mflists['types'] = fulltypelist
        mflists['values'] = fullvallist
        parset.update(mflists)
        
        # Rerun
        simparslist = parset.interp(start=project.settings.start, end=project.settings.end, dt=project.settings.dt)
        results = project.runsim(simpars=simparslist)
        updateplots(tmpresults=results, **kwargs)
        
    
    ## Keep the current parameters in the project; otherwise discard
    def keeppars():
        ''' Little function to reset origpars and update the project '''
        global origpars, tmppars, parset
        origpars = dcp(tmppars)
        parset.pars = tmppars
        project.parsets[name].pars = tmppars
        print('Parameters kept')
        return None
    
    
    def resetpars():
        ''' Reset the parameters to the last saved version -- WARNING, doesn't work '''
        global origpars, tmppars, parset
        tmppars = dcp(origpars)
        parset.pars = tmppars
        for i in range(nfull): boxes[i].setText(sigfig(fullvallist[i], sigfigs=nsigfigs))
        simparslist = parset.interp(start=project.settings.start, end=project.settings.end, dt=project.settings.dt)
        results = project.runsim(simpars=simparslist)
        updateplots(tmpresults=results)
        return None
    

    ## Set up GUI
    npars = len(fullkeylist)
    leftmargin = 10
    rowheight = 25
    colwidth = 450
    ncols = floor(npars/(maxrows+10*advanced))+1
    nrows = ceil(nfull/float(ncols))
    panelwidth = colwidth*ncols
    panelheight = rowheight*(nfull/ncols+2)+50
    buttonheight = panelheight-rowheight*1.5
    boxoffset = 300+leftmargin
    
    panel = pyqt.QWidget() # Create panel widget
    panel.setGeometry(100, 100, panelwidth, panelheight)
    spottaken = [] # Store list of existing entries, to avoid duplicates
    for i in range(nfull):
        row = (i % nrows) + 1
        col = floor(i/float(nrows))
        spot = (row,col)
        if spot in spottaken: 
            errormsg = 'Cannot add a button to %s since there already is one!' % str(spot)
            raise OptimaException(errormsg)
        else: spottaken.append(spot)
        
        texts.append(pyqt.QLabel(parent=panel))
        texts[-1].setText(fulllabellist[i])
        texts[-1].move(leftmargin+colwidth*col, rowheight*row)
        
        boxes.append(pyqt.QLineEdit(parent = panel)) # Actually create the text edit box
        boxes[-1].move(boxoffset+colwidth*col, rowheight*row)
        printv('Setting up GUI checkboxes: %s' % [i, fulllabellist[i], boxoffset+colwidth*col, rowheight*row], 4, verbose)
        boxes[-1].setText(sigfig(fullvallist[i], sigfigs=nsigfigs))
        boxes[-1].returnPressed.connect(manualupdate)
    
    keepbutton  = pyqt.QPushButton('Keep', parent=panel)
    resetbutton = pyqt.QPushButton('Reset', parent=panel)
    closebutton = pyqt.QPushButton('Close', parent=panel)
    
    keepbutton.move(1*panelwidth/4, buttonheight)
    resetbutton.move(2*panelwidth/4, buttonheight)
    closebutton.move(3*panelwidth/4, buttonheight)
    
    keepbutton.clicked.connect(keeppars)
    resetbutton.clicked.connect(resetpars)
    closebutton.clicked.connect(closewindows)
    panel.show()





def plotpeople(project=None, people=None, tvec=None, ind=None, simind=None, start=2, end=None, pops=None, animate=False, skipempty=True, verbose=2, toplot=None, **kwargs):
    '''
    A function to plot all people as a stacked plot
    
    "Exclude" excludes the first N health states -- useful for excluding susceptibles.
    
    Usage example:
        import optima as op
        P = op.defaults.defaultproject('simple')
        P.runsim()
        people = P.results[-1].raw[0]['people']
        op.plotpeople(P, people)
        
    NB: for a multiresult, simind must not be None!
    
    Version: 2018apr0
    '''
    if pops is None: pops = Ellipsis # This is a slice
    elif isnumber(pops): pops = [pops]
    if pops is not Ellipsis: plottitle = str(array(project.parsets[0].popkeys)[array(pops)])
    else: plottitle = 'All populations'
    legendsettings = {'loc':'upper left', 'bbox_to_anchor':(1.02, 1), 'fontsize':11, 'title':''}
    nocolor = (0.9,0.9,0.9)
    labels = project.settings.statelabels
    if toplot is None: toplot = 'people'
    
    if people is None:
        if ind is None: ind=-1
        try:
            people = project.results[ind].raw[0][toplot] # Try to get default people to plot
        except:
            if simind is None: simind = 1
            people = project.results[ind].raw[simind][0][toplot] # It's a multiresult: need another  index
        
    
    plotstyles = odict([
    ('susreg',   ('|','|')), 
    ('progcirc', ('+','|')), 
    ('undx',     ('O','o')), 
    ('dx',       ('.','o')), 
    ('care',     ('*','*')), 
    ('lost',     ('X','|')),
    ('usvl',     ('.','o')), 
    ('svl',      ('*','*')), 
    ])
    
    hatchstyles = []
    linestyles = []
    for key in plotstyles.keys():
        hatchstyles.extend([plotstyles[key][0] for lab in labels if lab.startswith(key)])
        linestyles.extend([plotstyles[key][1]  for lab in labels if lab.startswith(key)])
    
    labels = labels[start:end]
    hatchstyles = hatchstyles[start:end]
    linestyles = linestyles[start:end]
    
    ppl = people[start:end,:,:] # Exclude initial people
    ppl = ppl[:,pops,:] # Filter selected populations
    ppl = ppl[:,:,:].sum(axis=1) # Sum over people
    ppl = transpose(ppl) # So time is plotted on x-axis
    
    nstates = len(labels)
    colors = gridcolors(nstates)
    if tvec is None:
        tvec = project.settings.maketvec() # Won't necessarily match this ppl, supply as argument if so
    bottom = 0*tvec
    makenewfigure(**kwargs)
    ax = subplot(111)
    ylabel('Number of people')
    title(plottitle)
    xlim((tvec[0], tvec[-1]))
    for st in range(nstates-1,-1,-1):
        this = ppl[:,st]
        if sum(this): 
            thiscolor = colors[st]
            haspeople = True
        else: 
            thiscolor = nocolor
            haspeople = False
        if haspeople or not skipempty:
            printv('State: %i/%i Hatch: %s Line: %s Color: %s' % (st, nstates, hatchstyles[st], linestyles[st], thiscolor), 4, verbose)
            fill_between(tvec, bottom, this+bottom, facecolor=thiscolor, alpha=1, lw=0, hatch=hatchstyles[st])
            bottom += this
        
            # Legend stuff
            ax.plot((0, 0), (0, 0), color=thiscolor, linewidth=10, label=labels[st], marker=linestyles[st]) # This loop is JUST for the legends! since fill_between doesn't count as a plot object... -- TODO: this is copied from plotepi(), perhaps streamline
            handles, legendlabels = ax.get_legend_handles_labels()
            legend(reversed(handles), reversed(legendlabels), **legendsettings)
            if animate:
                show()
                pause(0.001)
    
    return None
    


global plotparsbackbut, plotparsnextbut, plotparslider
def plotpars(parslist=None, start=None, end=None, verbose=2, rows=6, cols=5, figsize=(16,12), fontsize=8, die=True, **kwargs):
    '''
    A function to plot all parameters. 'pars' can be an odict or a list of pars odicts.
    
    Version: 2016jan30
    '''
    from optima import Par, makesimpars, tic, toc
    from numpy import array, vstack
    import matplotlib.pyplot as plt
    from matplotlib.widgets import Button, Slider
    
    global position, plotparsbackbut, plotparsnextbut, plotparslider
    position = 0
    
    # In case the user tries to enter a project or parset -- TODO: make more flexible
    tmp = parslist
    try:  parslist = tmp.parsets[-1].pars # If it's a project
    except:
        try: parslist = tmp.pars # If it's a parset
        except: pass
    parslist = promotetolist(parslist) # Convert to list
    try:
        for i in range(len(parslist)): parslist[i] = parslist[i].pars
    except: pass # Assume it's in the correct form -- a list of pars odicts
    
    allplotdata = []
    for pars in parslist:
        count = 0
        simpars = makesimpars(pars, start=start, end=end)
        tvec = simpars['tvec']
        plotdata = array([['name','simpar','par_t', 'par_y']], dtype=object) # Set up array for holding plotting results
        for i,key1 in enumerate(pars):
            par = pars[key1]
            if isinstance(par, Par):
                if   hasattr(par,'y'): pardata = par.y # TODO: consider adding par.m as well
                elif hasattr(par,'p'): pardata = par.p # Population size
                else: raise Exception('???')
                if hasattr(pardata, 'keys') and len(pardata.keys())>0: # Only ones that don't have a len are temp pars
                    nkeys = len(pardata.keys())
                    for k,key2 in enumerate(pardata.keys()):
                        if hasattr(par, 't'): t = par.t[key2]
                        else: t = tvec[0] # For a constant
                        count += 1
                        if nkeys==1: thissimpar = simpars[key1]
                        else: thissimpar = simpars[key1][k]
                        thisplot = array(['%3i. %s - %s' % (count-1, key1, key2), thissimpar, t, pardata[key2]], dtype=object)
                        if array(thissimpar).sum()==0: thisplot[0] += ' (zero)'
                        plotdata = vstack([plotdata, thisplot])
                else:
                    t = tvec[0] # For a constant
                    count += 1
                    thisplot = array(['%3i. %s' % (count-1, key1), simpars[key1], t, pardata], dtype=object)
                    plotdata = vstack([plotdata, thisplot])
        plotdata = plotdata[1:,:] # Remove header
        allplotdata.append(plotdata)
    
    
    ## Do plotting
    nplots = len(plotdata)
    if any([len(pltd)!=nplots for pltd in allplotdata]): 
        printv('Warning, not all pars are the same length, only plotting first', 2, verbose)
        allplotdata = allplotdata[0]
    nperscreen = rows*cols

    plotparsfig = plt.figure(facecolor=(0.9,0.9,0.9), figsize=figsize)
    plt.subplots_adjust(left=0.05, right=0.95, bottom=0.1, top=0.95, wspace=0.3, hspace=0.4)
    plotparsaxs = []
    count = 0
    for row in range(rows):
        for col in range(cols):
            count += 1
            plotparsaxs.append(plotparsfig.add_subplot(rows, cols, count))
    
    backframe = plotparsfig.add_axes([0.1, 0.03, 0.1, 0.03])
    sliderframe = plotparsfig.add_axes([0.3, 0.03, 0.4, 0.03])
    nextframe = plotparsfig.add_axes([0.8, 0.03, 0.1, 0.03])
    plotparsbackbut = Button(backframe, 'Back')
    plotparsnextbut = Button(nextframe, 'Next')
    plotparslider = Slider(sliderframe, '', 0, nplots, valinit=0, valfmt='%d')
    
    def updateb(event=None): 
        global position
        position -= nperscreen
        position = max(0,position)
        position = min(nplots-nperscreen, position)
        plotparslider.set_val(position)
    
    def updaten(event=None): 
        global position
        position += nperscreen
        position = max(0,position)
        position = min(nplots-nperscreen, position)
        plotparslider.set_val(position)
    
    def update(tmp=0):
        global position, plotparslider
        position = tmp
        position = max(0,position)
        position = min(nplots-nperscreen, position)
        t = tic()
        for i,ax in enumerate(plotparsaxs):
            ax.cla()
            for item in ax.get_xticklabels() + ax.get_yticklabels(): item.set_fontsize(fontsize)
            ax.hold(True)
            nplt = i+position
            if nplt<nplots:
                for pd,plotdata in enumerate(allplotdata):
                    try:
                        this = plotdata[nplt,:]
                        ax.set_title(this[0])
                        if isinstance(this[1], dict):
                            if len(this[1].keys())==1:  this[1] = this[1][0] # Actually needs to be an odict
                            elif len(this[1].keys())>1: raise OptimaException('Expecting a number or an array or even an odict with one key, but got an odict with multiple keys (%s)' % this[0])
                        if   isnumber(this[1]):        ax.plot(tvec, 0*tvec+this[1])
                        elif len(this[1])==0:          ax.set_title(this[0]+' is empty')
                        elif len(this[1])==1:          ax.plot(tvec, 0*tvec+this[1])
                        elif len(this[1])==len(tvec):  ax.plot(tvec, this[1])
                        else: pass # Population size, doesn't use control points
                        printv('Plot %i/%i...' % (i*len(allplotdata)+pd+1, len(plotparsaxs)*len(allplotdata)), 2, verbose)
                    except Exception as E: 
                        if die: raise E
                        else: print('??????: %s' % repr(E))
                    try: 
                        if not(hasattr(this[3],'__len__') and len(this[3])==0): ax.scatter(this[2],this[3])
                    except Exception: pass # print('Problem with "%s": "%s"' % (this[0], repr(E)))
                    if pd==len(allplotdata)-1: # Do this for the last plot only
                        ax.set_ylim((0,1.1*ax.get_ylim()[1]))
                        ax.set_xlim((tvec[0],tvec[-1]))
        toc(t)
                
    update()
    plotparsbackbut.on_clicked(updateb)
    plotparsnextbut.on_clicked(updaten)
    plotparslider.on_changed(update)
    return allplotdata


def showplots(plots=None, figsize=None):
    '''
    This function can be used to show plots (in separate figure windows, independently
    of generating them.
    
    Example:
        import optima as op
        P = op.demo(0)
        plot = plotcascade(results=P.result(), interactive=False)
        op.showplots(plot) # Creates one plot
    
    NOTE: This function is purely remedial; the same effect can be accomplished more easily via:
        op.plotcascade(results=P.result(), interactive=True)

    Version: 2017may29
    '''
    ion()
    if figsize is None: figsize = (10,4)
    reanimateplots(plots) # Reconnect the plots to the matplotlib backend so they can be rendered
    nplots = len(plots)
    figs = []
    for p in range(nplots): 
        figs.append(figure(facecolor=(1,1,1),figsize=figsize))
        thisfig = figs[p]
        thisplot = plots[p].axes[0]
        thisfig._axstack.add(thisfig._make_key(thisplot), thisplot) # Add a plot to the axis stack
        thisplot.change_geometry(1, 1, 1) # Change geometry to be correct
        orig = thisplot.get_position() # get the original position 
        widthfactor = 0.9
        heightfactor = 0.9
        pos2 = [orig.x0, orig.y0,  orig.width*widthfactor, orig.height*heightfactor] 
        thisplot.set_position(pos2) # set a new position
    if nplots>1: return figs
    else:        return figs[0] # Don't return a list if a single figure


def loadplot(filename=None):
    '''
    Load a plot from a file and reanimate it.
    
    Example usage:
        import optima as op
        P = op.demo(0)
        op.saveplots(P, toplot='cascade', filetype='fig')
    
    Later:
        cascadefig = op.loadplot('cascade.fig')
    '''
    ion() # Without this, it doesn't show up
    fig = loadobj(filename)
    reanimateplots(fig)
    return fig



##############################################################################
### HELPER FUNCTIONS
##############################################################################

def makenewfigure(**figargs):
    ''' PyQt-specific function for maximizing the current figure '''
    global scrwid, scrhei
    
    if 'figsize' not in figargs: figargs['figsize'] = (scrwid, scrhei)
    if 'facecolor' not in figargs: figargs['facecolor'] = (1,1,1)
    fig = figure(**figargs) # Create a figure based on supplied kwargs, if any

    return fig
    

def closegui(event=None):
    ''' Close all GUI windows '''
    global panelfig, plotfig
    try: close(plotfig)
    except: pass
    try: close(panelfig)
    except: pass
    return None


def getchecked(check=None):
    ''' Return a list of whether or not each check box is checked or not '''
    ischecked = []
    for box in range(len(check.lines)): ischecked.append(check.lines[box][0].get_visible()) # Stupid way of figuring out if a box is ticked or not
    return ischecked


def clearselections(event=None):
    global check
    for box in range(len(check.lines)):
        for i in [0,1]: check.lines[box][i].set_visible(False)
    updateplots()
    return None


def defaultselections(event=None):
    ''' Reset to default options '''
    global check, results, globaladvanced
    plotselections = getplotselections(results, advanced=globaladvanced) # WARNING, assumes defaults don't change with advanced
    for box,tf in enumerate(plotselections['defaults']):
        if tf: # True if in defaults, false otherwise
            for i in [0,1]: check.lines[box][i].set_visible(True) # Two lines...stupid
        else:
            for i in [0,1]: check.lines[box][i].set_visible(False)
    updateplots()
    return None


def advancedselections(event=None):
    ''' Toggle advance doptions '''
    global check, checkboxes, updatebutton, clearbutton, defaultsbutton, advancedbutton, closebutton, plotfig, panelfig, results, plotargs, globaladvanced
    globaladvanced = not(globaladvanced) # Toggle
    try:    close(plotfig) # These work better here than caling closegui() directly
    except: pass
    try:    close(panelfig)
    except: pass
    check, checkboxes, updatebutton, clearbutton, defaultsbutton, advancedbutton, closebutton, plotfig, panelfig, plotargs = [None]*10 # Clear the bejesus out of everything
    pygui(results, advanced=globaladvanced)
    print('Switching to/from advanced; if GUI hangs, press enter in console') # Unfortunately, this happens from time to time
    pause(0.2) # Without this, it doesn't work...siiiigh
    return None
    

def zoomplots(event=None, ratio=1.0):
    ''' Zoom in or out '''
    global plotfig
    for ax in plotfig.axes:
        axpos = ax.get_position()
        x0 = axpos.x0
        x1 = axpos.x1
        y0 = axpos.y0
        y1 = axpos.y1
        xdiff = x1-x0
        ydiff = y1-y0
        xchange = xdiff*(1-ratio)/2.0
        ychange = ydiff*(1-ratio)/2.0
        ax.set_position([x0+xchange, y0+ychange, xdiff*ratio, ydiff*ratio])
    return None

def zoomin(event=None):
    ''' Zoom into plots '''
    zoomplots(event=event, ratio=1.1)
    return None

def zoomout(event=None):
    ''' Zoom out of plots '''
    zoomplots(event=event, ratio=0.9)
    return None

    
def updateplots(event=None, tmpresults=None, **kwargs):
    ''' Close current window if it exists and open a new one based on user selections '''
    global plotfig, check, checkboxes, results, plotargs
    if tmpresults is not None: results = tmpresults
    
    # If figure exists, get size, then close it
    if plotfig is None: plotfig = makenewfigure()
    width,height = plotfig.get_size_inches(); close(plotfig) # Get current figure dimensions
    
    # Get user selections
    ischecked = getchecked(check)
    toplot = array(checkboxes)[array(ischecked)].tolist() # Use logical indexing to get names to plot
    
    # Do plotting
    if sum(ischecked): # Don't do anything if no plots
        plotfig = makenewfigure(num='Optima results', figsize=(width, height), facecolor=(1,1,1))  # Create figure with correct number of plots
        for key in ['toplot','fig','figsize']: kwargs.pop(key, None) # Remove duplicated arguments if they exist
        plotresults(results, toplot=toplot, fig=plotfig, figsize=(width, height), **plotargs)
    
    return None





##############################################################################################################################
### GEOSPATIAL GUI
##############################################################################################################################


"""
GEOSPATIAL

This file defines everything needed for the Python GUI for geospatial analysis.

Version: 2017mar22
"""

from optima import Project, Portfolio, loadproj, saveobj,  defaultobjectives, makegeospreadsheet, makegeoprojects
from time import time

global geoguiwindow, globalportfolio, globalobjectives
if 1:  geoguiwindow, globalportfolio, globalobjectives = [None]*3


## Global options
budgetfactor = 1e6 # Conversion between screen and internal
prjext = '.prj'
prtext = '.prt'


##############################################################################################################################
## Define functions
##############################################################################################################################

def resetbudget():
    ''' Replace current displayed budget with default from portfolio '''
    global globalportfolio, objectiveinputs
    totalbudget = 0
    for project in globalportfolio.projects.values():
        totalbudget += sum(project.progsets[0].getdefaultbudget().values())
    objectiveinputs['budget'].setText(str(totalbudget/budgetfactor))
    return None


def warning(message, usegui=True):
    ''' usegui kwarg is so this can be used in a GUI and non-GUI context '''
    global geoguiwindow
    if usegui:
        pyqt.QMessageBox.warning(geoguiwindow, 'Message', message)
    else:
        print(message)
    
    
    
def gui_loadproj():
    ''' Helper function to load a project, since used more than once '''
    filepath = pyqt.QFileDialog.getOpenFileName(caption='Choose project file', filter='*'+prjext)
    project = None
    if filepath:
        try: project = loadproj(filepath, verbose=0)
        except Exception as E: print('Could not load file "%s": "%s"' % (filepath, repr(E)))
        if type(project)==Project: return project
        else: print('File "%s" is not an Optima project file' % filepath)
    else:
        print('No filepath provided')
    return project
    
    
def gui_makesheet():
    ''' Create a geospatial spreadsheet template based on a project file '''
    
    ## 1. Load a project file
    project = gui_loadproj() # No, it's a project path, load it
    if project is None: 
        raise OptimaException('No project loaded.')
    
    try:    results = project.parsets[-1].getresults()
    except: results = project.runsim(name=project.parsets[-1].name)
    
    copies, ok = pyqt.QInputDialog.getText(geoguiwindow, 'GA Spreadsheet Parameter', 'How many variants of the chosen project do you want?')
    try: copies = int(copies)
    except: raise OptimaException('Input (number of project copies) cannot be converted into an integer.')
    
    refyear, ok = pyqt.QInputDialog.getText(geoguiwindow, 'GA Spreadsheet Parameter', 'Select a reference year for which you have district data.')
    try: refyear = int(refyear)
    except: raise OptimaException('Input (reference year) cannot be converted into an integer.')
    if not refyear in [int(x) for x in results.tvec]:
        raise OptimaException("Input not within range of years used by aggregate project's last stored calibration.")

    ## 2. Get destination filename
    spreadsheetpath = pyqt.QFileDialog.getSaveFileName(caption='Save geospatial spreadsheet file', filter='*.xlsx')
    
    # 4. Generate and save spreadsheet
    try:
        makegeospreadsheet(project=project, filename=spreadsheetpath, copies=copies, refyear=refyear, verbose=2)
        warning('Multi-project template saved to "%s".' % spreadsheetpath)
    except:
        warning('Error: Template not saved due to a workbook error!')

    return None
    
    
def gui_makeproj():
    ''' Create a series of project files based on a seed file and a geospatial spreadsheet '''
    project = gui_loadproj()
    spreadsheetpath = pyqt.QFileDialog.getOpenFileName(caption='Choose geospatial spreadsheet', filter='*.xlsx')
    destination = pyqt.QFileDialog.getExistingDirectory(caption='Choose output folder')
    makegeoprojects(project=project, spreadsheetpath=spreadsheetpath, destination=destination)
    warning('Created projects from spreadsheet')
    return None


def gui_create(filepaths=None, portfolio=None, doadd=False):
    ''' Create a portfolio by selecting a list of projects; silently skip files that fail '''
    global globalportfolio, projectslistbox, objectiveinputs
    
    projectpaths = []
    projectslist = []
    if globalportfolio is None: 
        globalportfolio = Portfolio()
    if not doadd:
        globalportfolio = Portfolio()
        projectslistbox.clear()
    if doadd and portfolio != None:
        globalportfolio = portfolio
    filepaths = pyqt.QFileDialog.getOpenFileNames(caption='Choose project files', filter='*'+prjext)
    if filepaths:
        if type(filepaths)==str: filepaths = [filepaths] # Convert to list
        for filepath in filepaths:
            tmpproj = None
            try: tmpproj = loadproj(filepath, verbose=0)
            except: print('Could not load file "%s"; moving on...' % filepath)
            if tmpproj is not None: 
                if type(tmpproj)==Project:
                    projectslist.append(tmpproj)
                    projectpaths.append(filepath)
                    print('Project file "%s" loaded' % filepath)
                else: print('File "%s" is not an Optima project file; moving on...' % filepath)
        projectslistbox.addItems(projectpaths)
        globalportfolio.addprojects(projectslist)
        resetbudget() # And reset the budget
    return None


def gui_addproj():
    ''' Add a project -- same as creating a portfolio except don't overwrite '''
    gui_create(doadd=True)
    resetbudget() # And reset the budget
    return None


def gui_loadport():
    ''' Load an existing portfolio '''
    global globalportfolio, projectslistbox
    filepath = pyqt.QFileDialog.getOpenFileName(caption='Choose portfolio file', filter='*'+prtext)
    tmpport = None
    if filepath:
        try: tmpport = loadobj(filepath, verbose=0)
        except Exception as E: 
            warning('Could not load file "%s" because "%s"' % (filepath, repr(E)))
            return None
        if tmpport is not None: 
            if type(tmpport)==Portfolio:
                globalportfolio = tmpport
                projectslistbox.clear()
                projectslistbox.addItems([proj.name for proj in globalportfolio.projects.values()])
                print('Portfolio file "%s" loaded' % filepath)
            else: print('File "%s" is not an Optima portfolio file' % filepath)
    else:
        warning('File path not provided. Portfolio not loaded.')
    resetbudget() # And reset the budget
    return None


def gui_rungeo():
    ''' Actually run geospatial analysis!!! '''
    global globalportfolio, globalobjectives, objectiveinputs
    starttime = time()
    if globalobjectives is None:
        globalobjectives = defaultobjectives()
        globalobjectives['budget'] = 0.0 # Reset
    for key in objectiveinputs.keys():
        globalobjectives[key] = eval(str(objectiveinputs[key].text())) # Get user-entered values
    globalobjectives['budget'] *= budgetfactor # Convert back to internal representation
    BOCobjectives = dcp(globalobjectives)
    try:
        globalportfolio.genBOCs(objectives=BOCobjectives, maxtime=30, mc=0)
        globalportfolio.runGA(objectives=globalobjectives, maxtime=30, reoptimize=True, mc=0, batch=True, verbose=2, die=False, strict=True)
    except Exception as E:
        warning('Geospatial analysis failed: %s' % repr(E))
    warning('Geospatial analysis finished running; total time: %0.0f s' % (time() - starttime))
    return None
    

def gui_plotgeo():
    ''' Actually plot geospatial analysis!!! '''
    global globalportfolio
    if globalportfolio is None: 
        warning('Please load a portfolio first')
        return None
    globalportfolio.plotBOCs(deriv=False)
    return None


def gui_export():
    ''' Save the current results to Excel file '''
    global globalportfolio
    if type(globalportfolio)!=Portfolio: warning('Warning, must load portfolio first!')
    
    # 2. Create a new file dialog to save this spreadsheet
    filepath = pyqt.QFileDialog.getSaveFileName(caption='Save geospatial analysis results file', filter='*.xlsx')
    
    # 3. Generate spreadsheet according to David's template to store these data
    if filepath:
        try:
            globalportfolio.export(filename=filepath)
        except Exception as E:
            warning('Results export failed: %s' % repr(E))
        warning('Results saved to "%s".' % filepath)
    else:
        warning('Filepath not supplied: %s' % filepath)
    
    return None
    

def gui_saveport():
    ''' Save the current portfolio '''
    global globalportfolio
    filepath = pyqt.QFileDialog.getSaveFileName(caption='Save portfolio file', filter='*'+prtext)
    saveobj(filepath, globalportfolio)
    return None


def closewindow(): 
    ''' Close the control panel '''
    global geoguiwindow
    geoguiwindow.close()
    return None
    

def geogui():
    '''
    Open the GUI for doing geospatial analysis.
    
    Version: 2016jan23
    '''
    global geoguiwindow, globalportfolio, globalobjectives, objectiveinputs, projectslistbox, projectinfobox
    globalportfolio = None
    if globalobjectives is None:
        globalobjectives = defaultobjectives()
        globalobjectives['budget'] = 0.0 # Reset
    
    ## Set parameters
    wid = 1200.0
    hei = 600.0
    top = 20
    spacing = 40
    left = 20.
    
    ## Housekeeping
    fig = figure(); close(fig) # Open and close figure...dumb, no? Otherwise get "QWidget: Must construct a QApplication before a QPaintDevice"
    geoguiwindow = pyqt.QWidget() # Create panel widget
    geoguiwindow.setGeometry(100, 100, wid, hei)
    geoguiwindow.setWindowTitle('Optima geospatial analysis')
    
    ##############################################################################################################################
    ## Define buttons
    ##############################################################################################################################
    
    ## Define buttons
    buttons = odict()
    buttons['makesheet'] = pyqt.QPushButton('Make geospatial spreadsheet from project', parent=geoguiwindow)
    buttons['makeproj']  = pyqt.QPushButton('Auto-generate projects from spreadsheet', parent=geoguiwindow)
    buttons['create']    = pyqt.QPushButton('Create portfolio from projects', parent=geoguiwindow)
    buttons['add']       = pyqt.QPushButton('Add projects to portfolio', parent=geoguiwindow)
    buttons['loadport']  = pyqt.QPushButton('Load existing portfolio', parent=geoguiwindow)
    buttons['rungeo']    = pyqt.QPushButton('Run geospatial analysis', parent=geoguiwindow)
    buttons['plotgeo']   = pyqt.QPushButton('Plot geospatial results', parent=geoguiwindow)
    buttons['export']    = pyqt.QPushButton('Export results', parent=geoguiwindow)
    buttons['saveport']  = pyqt.QPushButton('Save portfolio', parent=geoguiwindow)
    buttons['close']     = pyqt.QPushButton('Close', parent=geoguiwindow)
    
    ## Define button functions
    actions = odict()
    actions['makesheet'] = gui_makesheet
    actions['makeproj']  = gui_makeproj
    actions['create']    = gui_create
    actions['add']       = gui_addproj
    actions['loadport']  = gui_loadport
    actions['rungeo']    = gui_rungeo
    actions['plotgeo']   = gui_plotgeo
    actions['export']    = gui_export
    actions['saveport']  = gui_saveport
    actions['close']     = closewindow
    
    ## Set button locations
    spacer = 0
    for b,key in enumerate(buttons.keys()):
        if key=='rungeo': spacer = 170
        buttons[key].move(left, top+spacing*b+spacer)
    
    ## Define button functions
    for key in buttons.keys():
        buttons[key].clicked.connect(actions[key])
    
    
    
    ##############################################################################################################################
    ## Define other objects
    ##############################################################################################################################
    
    def updateprojectinfo():
        global globalportfolio, projectslistbox, projectinfobox
        ind = projectslistbox.currentRow()
        project = globalportfolio.projects[ind]
        projectinfobox.setText(repr(project))
        return None
    
    def removeproject():
        global projectslistbox, projectinfobox, globalportfolio
        ind = projectslistbox.currentRow()
        globalportfolio.projects.pop(globalportfolio.projects.keys()[ind]) # Remove from portfolio
        projectslistbox.takeItem(ind) # Remove from list
        return None
        
    
    ## List of projects
    projectslistlabel = pyqt.QLabel(parent=geoguiwindow)
    projectslistlabel.setText('Projects in this portfolio:')
    projectslistbox = pyqt.QListWidget(parent=geoguiwindow)
    projectslistbox.verticalScrollBar()
    projectslistbox.currentItemChanged.connect(updateprojectinfo)
    buttons['remove'] = pyqt.QPushButton('Remove selected project from portfolio', parent=geoguiwindow)
    buttons['remove'].clicked.connect(removeproject)
    projectslistlabel.move(330,20)
    projectslistbox.move(330, 40)
    buttons['remove'].move(330, hei-40)
    projectslistbox.resize(300, hei-100)
    
    
    ## Project info
    projectsinfolabel = pyqt.QLabel(parent=geoguiwindow)
    projectsinfolabel.setText('Information about the selected project:')
    projectinfobox = pyqt.QTextEdit(parent=geoguiwindow)
    projectinfobox.setReadOnly(True)
    projectinfobox.verticalScrollBar()
    projectsinfolabel.move(640,20)
    projectinfobox.move(640, 40)
    projectinfobox.resize(530, hei-100)
    
    ## Objectives
    objectivetext = odict()
    objectivetext['start']       = 'Start year:'
    objectivetext['end']         = 'End year:'
    objectivetext['budget']      = 'Total budget (mil.):'
    objectivetext['deathweight'] = 'Deaths weight:'
    objectivetext['inciweight']  = 'Infections weight:'
    
    objectivetextobjs = odict()
    for k,key in enumerate(objectivetext.keys()):
        objectivetextobjs[key] = pyqt.QLabel(parent=geoguiwindow)
        objectivetextobjs[key].setText(str(objectivetext[key]))
        objectivetextobjs[key].move(left+10, 235+k*30)
    
    objectiveinputs = odict()
    for k,key in enumerate(objectivetext.keys()):
        objectiveinputs[key] = pyqt.QLineEdit(parent=geoguiwindow)
        objectiveinputs[key].setText(str(globalobjectives[key]))
        objectiveinputs[key].move(left+120, 230+k*30)
    objectiveinputs['budget'].setText(str(globalobjectives['budget']/budgetfactor)) # So right units
    

    geoguiwindow.show()