import os, stat, shutil

import yoda
from yoda.plotting.mlp_preprocessor import preprocess
import yoda.plotting.fetch_data as fd

import matplotlib as mpl
import matplotlib.pyplot as plt
mpl.use('Agg')

import math, re
import numpy as np


def setCanvas(cmd, nRatios):
    """Return Pythonic string corresponding to
    basic mpl canvas, with optional splitting."""
    ## Create fig and axes
    cmd.newline()
    if nRatios == 0:
        return cmd.add("# Create figure and axis objects", "fig, ax = plt.subplots(1, 1)")
    if nRatios > 1:
        cmd.add("# Adjust figure aspect ratio for number of ratio panels.",
                "# (height of main panel assumed twice that of ratio panel)",
                "nRatioPanels = %d" % nRatios,
                "figH *= (nRatioPanels+2) / 3.0")
        cmd.newline()
    cmd.add("# Create figure and axis objects")
    raxes = ', '.join([ 'ratio%d_ax' % i for i in range(nRatios) ])
    rsplits = ', '.join(['1']*nRatios)
    cmd.add(f"fig, (ax, {raxes}) = plt.subplots({nRatios+1:d}, 1, sharex=True,")
    cmd.add(f"                  figsize=(figW,figH), gridspec_kw={{'height_ratios': (2, {rsplits})}})")


def setStyleFile(cmd, yaml_dicts, outdir):

    ## Set name of style file
    mpl_stylename = yaml_dicts.get('style', 'default') + '.mplstyle'

    ## Find location of style file
    plot_style = None
    for styledir in yoda.getYodaDataPath():
        if not os.path.exists(styledir):
            continue
        if os.path.exists( os.path.join(styledir, 'plotting/') ):
            ## Check if the style file is in the standard share/YODA/plotting dir
            if os.path.isfile( os.path.join(styledir, 'plotting/'+ mpl_stylename) ):
                plot_style = os.path.join(styledir, 'plotting/' + mpl_stylename)
        elif os.path.isfile( os.path.join(styledir, mpl_stylename) ):
            ## Check if it's in a user-supplied dir
            plot_style = os.path.join(styledir, mpl_stylename)
    if plot_style is None:
        raise NotImplementedError(f'Plot style {mpl_stylename} file not found. Make sure it' \
                                    ' can be found in one of these directories:' \
                                    f'{yoda.getYodaDataPath()}')

    ## Put copy of style file into output directory
    if not os.path.isfile(os.path.join(outdir, mpl_stylename)):
        os.system(f"cp {plot_style} {os.path.join(outdir, mpl_stylename)}")

    ## Tell mpl about style
    plt.style.use(plot_style) # why is this required for the generating script?
    stylepath = yaml_dicts.get('stylepath', '') + mpl_stylename
    cmd.add("#plot style", f"plt.style.use(os.path.join(plotDir, '{stylepath}'))")



# TODO: There seems to be a lot of overlap with mkPlottingScript2D -> streamline
def mkPlottingScript1D(hist_data, hist_features, yaml_dicts, outdir, plot_name,
                       multiweight_dict, formats, has_latex = True):
    """Returns Pythonic string corresponding to standalone plotting script."""

    cmd = yoda.util.StringCommand()
    setStyleFile(cmd, yaml_dicts, outdir)
    plot_features = yaml_dicts.get('plot features', {})

    if not has_latex:
        cmd.newline()
        cmd.add(f"plt.rcParams['text.usetex'] = False",
                f"plt.style.use(['classic'])",
                f"plt.rc('font', family='DejaVu Sans')")

    cmd.add("# plot metadata",
            f"""figW, figH = plt.rcParams['figure.figsize']""",
            f"""ax_xLabel = r'{plot_features.get("XLabel", "")}'""",
            f"""ax_yLabel = r'{plot_features.get("YLabel", "")}'""",
            f"""ax_title  = r'{plot_features.get("Title", "")}'""")

    xscale = 'log' if int(plot_features.get('LogX', 0)) else 'linear'
    yscale = 'log' if int(plot_features.get('LogY', 1)) else 'linear'
    cmd.add(f"""ax_xScale = '{xscale}'""",
            f"""ax_yScale = '{yscale}'""")

    ax_format = {}  # Stores the items for formatting the axes in a dict

    deviation = plot_features.get('Deviation')

    ## Set plot limits
    xmin = plot_features.get('XMin', min([h.xMin() for h in hist_data.values()]))
    xmax = plot_features.get('XMax', max([h.xMax() for h in hist_data.values()]))
    ax_format['xlim'] = (float(xmin), float(xmax))

    ## Set maximum y-value from all hist datasets
    #max_ymax = max([h.points()[0].val(1) for h in hist_data.values()])
    max_ymax = max([max(h.yVals()) for h in hist_data.values()])
    if plot_features.get('YMax') is not None:
        ymax = plot_features.get('YMax')
    elif plot_features.get('LogY', 1):
        ## Round off highest number in the histograms to next power of 10
        ymax = 1.1 * 10**(math.ceil(math.log10(max_ymax))) if max_ymax > 0 else 1.0
    else:
        ymax = 1.1*max_ymax

    ## Use minimum y-value from all hist datasets
    #min_ymin = min([h.points()[0].val(1) for h in hist_data.values()])
    #min_ymin_positive = min([h.points()[0].val(1) for h in hist_data.values() if h.points()[0].val(1) > 0 ])
    min_ymin = min([min(h.yVals()) for h in hist_data.values()])
    min_ymin_pos_l = []
    for h in hist_data.values():
        min_ymin_pos_l += [y for y in h.yVals() if y > 0]
    min_ymin_positive = min(min_ymin_pos_l) if len(min_ymin_pos_l) else 0 # len==0 when all histo values are 0

    if plot_features.get('YMin') is not None:
        ymin = plot_features.get('YMin')
    elif plot_features.get('LogY', 1):
        ## Round off lowest number in the histograms to lower power of 10
        # TODO: come up with a better solution to deal with min_ymin=0
        ymin = 0.9 * 10**(math.floor(math.log10(min_ymin_positive))) if min_ymin_positive != 0 else 1e-3*ymax
    elif plot_features.get('ShowZero', 1):  # default ShowZero is True
        ymin = 0 if min_ymin > -1e-4 else 1.1*min_ymin
    else:
        ymin = (1.1*min_ymin if min_ymin < -1e-4 else 0 if min_ymin < 1e-4 else 0.9*min_ymin)

    ax_format['ylim'] = (float(ymin), float(ymax))
    ax_format['logx'] = int(plot_features.get('LogX', 0))
    ax_format['logy'] = int(plot_features.get('LogY', 1))

    if ax_format.get('xlim'): cmd.add(f"""xLims = {ax_format['xlim']}""")
    if ax_format.get('ylim'): cmd.add(f"""yLims = {ax_format['ylim']}""")

    ## Labels for the legend
    labels = []
    if plot_features.get('Legend', 1):
        for i in range(len(hist_data.values())):
            lentry = hist_features[i].get('Title', 'Curve {}'.format(i+1))
            if lentry is not None and lentry.count('_') and lentry.count('$') < 2:
                lentry = lentry.replace(".yoda","").replace('_', r'\_')
            if lentry is not None:
                labels.append(lentry)
        custom_legend = plot_features.get('CustomLegend', '')
        if custom_legend:
            labels.append(custom_legend)
    if len(labels) > 0:
      texLabels = 'r"' + '", r"'.join(labels) + '"' #< We use r-strings to be TeX friendly
      cmd.add(f"# TeX-friendly labels for the legend", f"labels = [ {texLabels} ]")


    ratio_panels = dict()
    for label in plot_features.keys():
      ## Unless specifically deactivated, default is to always have at least one ratio panel
      default = label == 'RatioPlot' and plot_features.get('RatioPlot', 1)
      if default or re.fullmatch("RatioPlot(\d+)?", label) is not None:
          ratio_ref = plot_features.get(label + 'DrawReference', None)
          if ratio_ref is not None:
              ratio_panels[label] = [ ratio_ref ]
          elif default:
              ratio_panels[label] = [ list(hist_data.keys())[0] ]
          ratio_lines = plot_features.get(label + 'DrawOnly', None)
          if ratio_lines is not None:
              for rcurve in ratio_lines.split():
                  ratio_panels.append( rcurve )
          elif default:
              ratio_panels[label] += list(hist_data.keys())[1:]
          n = '0' if 't' == label[-1] else label[-1]
          ax_format['ratio%s_logy' % n] = plot_features.get(label + 'LogY', 0)
    ## Unless user specifically asks not have ratio panel, have one by default
    if not ratio_panels and plot_features.get('RatioPlot', 1):
        ratio_panels['RatioPlot'] = list(hist_data.keys())
        ax_format['ratio0_logy'] = plot_features.get('RatioPlotLogY', 0)

    ## Split canvas into panels
    setCanvas(cmd, len(ratio_panels.keys()))

    ## Set some ratio-panel formatting defaults
    for i, key in enumerate(ratio_panels.keys()):
        ratioName = 'ratio%d_ax' % i
        ## Ratio plot has default y-range of 0.5 to 1.5
        if deviation:
            ratioYMin = -3.0
            ratioYMax =  3.0
            #ratioYLabel = 'Ratio'
            cmd.add(f"{ratioName}.set_yticks(list(range(-3,3)),")
            cmd.add(f"    [r'$-3\,\sigma$', r'$-2\,\sigma$', r'$-1\,\sigma$', r'$0\,\sigma$', r'$1\,\sigma$', r'$2\,\sigma$']  )")
        else:
            ratioYMin = plot_features.get(key+'YMin', 0.5)
            ratioYMax = plot_features.get(key+'YMax', 1.4999)  # avoid 1.5 tick mark
            cmd.add(f"{ratioName}.yaxis.set_major_locator(mpl.ticker.MultipleLocator(0.1))")
        ratioYLabel = plot_features.get(key+'YLabel', 'Ratio')
        cmd.add(f"{ratioName}.set_ylabel('{ratioYLabel}')",
                f"{ratioName}.set_ylim({ratioYMin}, {ratioYMax})")
        if plot_features.get('DummyXaxis'):
            cmd.add(f"{ratioName}.tick_params(axis='x', which='both', bottom=False, top=False, labelbottom=False)")

    ax_format['xmajor_ticks'] = plot_features.get('XMajorTickMarks')
    ax_format['ymajor_ticks'] = plot_features.get('YMajorTickMarks')
    ax_format['xminor_ticks'] = plot_features.get('XMinorTickMarks')
    ax_format['yminor_ticks'] = plot_features.get('YMinorTickMarks')
    ax_format['xcustom_major_ticks'] = plot_features.get('XMajorTickMarks')

    if plot_features.get('XCustomMajorTicks') is not None:
        ax_format['xcustom_major_ticks'] = plot_features.get('XMajorTickMarks')
        for i in range(len(ratio_panels.keys())):
            ratioName = 'ratio%d_ax' % i
            cmd.add(f"{ratioName}.set_xticks([], minor=True)")

    ax_format['ycustom_major_ticks'] = plot_features.get('YMajorTickMarks')
    ax_format['ycustom_minor_ticks'] = plot_features.get('YMajorTickMarks')
    ax_format['ycustom_minor_ticks'] = plot_features.get('YMajorTickMarks')
    ax_format['plot_xticklabels'] = plot_features.get('PlotXTickLabels')
    for i, key in enumerate(ratio_panels.keys()):
        ratioName = 'ratio%d_ax_' % i
        ax_format[ratioName + 'ycustom_major_ticks'] = plot_features.get(ratioName+'YMajorTickMarks')
        ax_format[ratioName + 'ycustom_minor_ticks'] = plot_features.get(ratioName+'YMajorTickMarks')
        ax_format[ratioName + 'ycustom_minor_ticks'] = plot_features.get(ratioName+'YMajorTickMarks')

    plot_errorbars = [h.get('ErrorBars', 1) and not h.get('ErrorBands', 0) for h in hist_features]
    default_colors = plt.rcParams['axes.prop_cycle'].by_key()['color']
    col_idx = 0; colors = [ ]
    for i, (label,ao) in enumerate(hist_data.items()):
        # move command-line tags into annotations
        for tag in ['IsRef', 'LineStyle', 'MarkerStyle', 'ErrorBandOpacity']:
            if tag in hist_features[i]:
                ao.setAnnotation(tag, hist_features[i][tag])
                del hist_features[i][tag]
        if 'LineColor' in hist_features[i]:
            colors.append( hist_features[i]['LineColor'] )
            del hist_features[i]['LineColor']
        else:
            colors.append( default_colors[col_idx % len(default_colors)] )
            col_idx += 1
    cmd.newline()
    cmd.add("# LineColor map: curve index -> color",
            f"colors = {str({i: col for i, col in enumerate(colors)})}")

    ## Write out the plot data to __data.py file
    dataOutPyName = plot_name.strip('/') + '__data.py'
    dataOutPyImport = plot_name.split('/')[-1] + '__data.py'

    cmd.add(f"# the numerical data is stored in a separate file",
            f"dataf = dict()",
            f"prefix = os.path.split(__file__)[0]",
            f"if prefix:    prefix = prefix + '/'",
            f"exec(open(prefix+'{dataOutPyImport}').read(), dataf)")

    # propagate IsRef info that was passed to hist_features instead of already being an annotation in the input YODA file
    for i, (k,v) in enumerate(hist_data.items()):
        if 'IsRef' in hist_features[i]:
            v.setAnnotation('IsRef', hist_features[i]['IsRef'])

    with open(os.path.join(outdir, dataOutPyName), "a") as dataOutFile:
        ## Fetch numerical data for plotting from separate function
        datacmd, data = fd.mkCurves1D(hist_data, ratio_panels, plot_errorbars, multiweight_dict, colors, deviation)
        cmd.add(datacmd)
        dataOutFile.write(data)

    ## Add a dummy handle if there's a custom legend entry
    if plot_features.get('CustomLegend', ''):
        cmd.add("import matplotlib.patches as mpatches",
                "legend_handles.append(mpatches.Patch(color='none'))")
    ## Fill, position and draw legend
    if plot_features.get('Legend', 1):
        if plot_features.get('LegendAlign') is None or plot_features.get('LegendAlign') == 'l':
            legend_pos = (float(plot_features.get('LegendXPos', 0.5)),
                          float(plot_features.get('LegendYPos', 0.97)))
            cmd.add(f"legend_pos = {legend_pos}",
                    f"ax.legend(legend_handles, labels, loc='upper left', bbox_to_anchor=legend_pos)")
        if plot_features.get('LegendAlign') == 'r':
            legend_pos = (float(plot_features.get('LegendXPos', 0.97)),
                          float(plot_features.get('LegendYPos', 0.97)))

            cmd.add(f"legend_pos = {legend_pos}",
                    f"ax.legend(legend_handles, labels, loc='upper right',",
                    f"bbox_to_anchor=legend_pos,markerfirst=False)")

    ## Set text labels on axes
    cmd.newline()
    cmd.add("# set plot metadata as defined above")
    for i in range(len(ratio_panels.keys())):
        ratioName = 'ratio%d_ax' % i
        cmd.add(f"{ratioName}.set_xlabel(ax_xLabel)")
    if not len(ratio_panels.keys()):
        cmd.add(f"ax.set_xlabel(ax_xLabel)")

    cmd.add(f"ax.set_ylabel(ax_yLabel, loc='top')",
            f"ax.set_title(ax_title, loc='left')",
            f"ax.set_xscale(ax_xScale)",
            f"ax.set_yscale(ax_yScale)")

    ## Toggle x/y lims
    if ax_format['xlim']: cmd.add("ax.set_xlim(xLims)")
    if ax_format['ylim']: cmd.add("ax.set_ylim(yLims)")

    cmd.newline()
    cmd.add(f"# tick formatting",
            f"plt.rcParams['xtick.top'] = {plot_features.get('XTwosidedTicks', True)}",
            f"plt.rcParams['ytick.right'] = {plot_features.get('YTwosidedTicks', True)}")

    # TODO: What about setting log scale for the ratio panels?
    if ax_format['logy']:
        cmd.add(f"ax.yaxis.set_major_locator(mpl.ticker.LogLocator(base=10.0, numticks=np.inf))",
                f"ax.yaxis.set_minor_locator(mpl.ticker.LogLocator(",
                f"                           base=10.0, subs=np.arange(0.1, 1, 0.1), numticks=np.inf))")
    for i in range(len(ratio_panels.keys())):
        ratioName = 'ratio%d' % i
        if ax_format[ratioName + '_logy']:
            cmd.add(f"{ratioName}_ax.yaxis.set_major_locator(mpl.ticker.LogLocator(base=10.0, numticks=np.inf))",
                    f"{ratioName}_ax.yaxis.set_minor_locator(mpl.ticker.LogLocator(",
                    f"                           base=10.0, subs=np.arange(0.1, 1, 0.1), numticks=np.inf))")
    if plot_features.get('DummyXaxis'):
        cmd.add("ax.tick_params(axis='x', which='both', bottom=False, top=False, labelbottom=False)")

    if plot_features.get('canvasText', None):
        ## Escape apostrophes
        text = plot_features.get('canvasText').replace("'", r"\'")
        cmd.add(f"ax.text(0.05, 0.95, '{text}', ha='left', va='top', transform=ax.transAxes)")

    if len(ratio_panels.keys()):
        rpanels = ', '.join([ 'ratio%d_ax' % i for i in range(len(ratio_panels.keys())) ])
        cmd.add(f"fig.align_ylabels((ax, {rpanels}))")

    # https://stackoverflow.com/questions/8213522/when-to-use-cla-clf-or-close-for-clearing-a-plot-in-matplotlib
    # plt.close(fig) instead of fig.clf()?
    for fmt in formats:
        figName = plot_name.split('/')[-1]+f'.{fmt.lower()}'
        cmd.add(f"plt.savefig(os.path.join(plot_outdir, '{figName}'), format='{fmt}')")
    cmd.newline()
    cmd.add(f"plt.close(fig)")

    return cmd.get()


def mkPlottingScript2D(hist_data, hist_features, yaml_dicts, outdir,
                       plot_name, formats, has_latex = True):
    """ Return Python commands to make a Rivet plot from the input parameters."""

      # TODO: lots of code duplication with 1D plotting: we should try to unify

    cmd = yoda.util.StringCommand()
    setStyleFile(cmd, yaml_dicts, outdir)
    plot_features = yaml_dicts.get('plot features', {})

    if not has_latex:
        cmd.newline()
        cmd.add(f"plt.rcParams['text.usetex'] = False",
                f"plt.style.use(['classic'])",
                f"plt.rc('font', family='DejaVu Sans')")

    cmd.add(f"# plot metadata",
            f"""figW, figH = plt.rcParams['figure.figsize']""",
            f"ax_xLabel = r'{plot_features.get('XLabel', '')}'",
            f"ax_yLabel = r'{plot_features.get('YLabel', '')}'",
            f"ax_zLabel = r'{plot_features.get('ZLabel', '')}'",
            f"ax_title  = r'{plot_features.get('Title', '')}'")

    xscale = 'log' if int(plot_features.get('LogX', 0)) else 'linear'
    yscale = 'log' if int(plot_features.get('LogY', 1)) else 'linear'
    zscale = 'log' if int(plot_features.get('LogZ', 0)) else 'linear'
    cmd.add(f"ax_xScale = '{xscale}'")
    cmd.add(f"ax_yScale = '{yscale}'")
    cmd.add(f"ax_zScale = '{zscale}'")

    ax_format = {}  # Stores the items for formatting the axes in a dict
    # Set plot limits

    ## Set maximum Y value from all hist datasets
    max_ymax = max([h.points()[0].val(1) for h in hist_data])

    if plot_features.get('YMax') is not None:
        ymax = plot_features.get('YMax')
    elif plot_features.get('LogY', 1):
        ## Round off the highest number in the histograms to the next power of 10
        ymax = 10**(math.ceil(math.log10(max_ymax)))
    else:
        ymax = 1.1*max_ymax

    ## Use minimum y-value from all hist datasets
    min_ymin = min([h.points()[0].val(1) for h in hist_data]) # TO DO -- where does this come from??
    if plot_features.get('YMin') is not None:
        ymin = plot_features.get('YMin')
    elif plot_features.get('LogY', 1):
        ## Round off lowest number in the histograms to lower power of 10
        # TODO: come up with a better solution to deal with min_ymin=0
        ymin = 10**(math.floor(math.log10(min_ymin))) if min_ymin > 0 else 2e-7*ymax
    elif plot_features.get('ShowZero', 1):  # default ShowZero is True
        ymin = 0 if min_ymin > -1e-4 else 1.1*min_ymin
    else:
        ymin = (1.1*min_ymin if min_ymin < -1e-4 else 0 if min_ymin < 1e-4
                else 0.9*min_ymin)

    ax_format['ylim'] = (ymin, ymax)
    ax_format['logx'] = int(plot_features.get('LogX', 0))
    ax_format['logy'] = int(plot_features.get('LogY', 1))
    ax_format['logz'] = int(plot_features.get('LogZ', 0))

    if ax_format.get('xlim'):  cmd.add(f"xLims = {ax_format['xlim']}")
    if ax_format.get('ylim'):  cmd.add(f"yLims = {ax_format['ylim']}")

    ## Labels for the legend
    labels = []
    if plot_features.get('Legend', 1):
        for i in range(len(hist_data)):
            lentry = hist_features[i].get('Title', 'Curve {}'.format(i+1))
            if lentry.count('_') and lentry.count('$') < 2:
                lentry.replace(".yoda","").replace('_', r'\_')
            labels.append(lentry)
    if len(labels) > 0:
      texLabels = 'r"' + '", r"'.join(labels) + '"' ## We use r-strings to be TeX-friendly
      cmd.add(f"# TeX-friendly labels for the legend", f"labels = [ {texLabels} ]")

    setCanvas(cmd, 0) # no ratio panel

    if ax_format['logy']:
      cmd.add(f"ax.yaxis.set_minor_locator(mpl.ticker.LogLocator(",
              f"base=10.0, subs=[i for i in np.arange(0, 1, 0.1)], numticks=np.inf))")
    ax_format['xmajor_ticks'] = plot_features.get('XMajorTickMarks')
    ax_format['ymajor_ticks'] = plot_features.get('YMajorTickMarks')
    ax_format['xminor_ticks'] = plot_features.get('XMinorTickMarks')
    ax_format['yminor_ticks'] = plot_features.get('YMinorTickMarks')
    ax_format['xcustom_major_ticks'] = plot_features.get('XMajorTickMarks')

    if plot_features.get('XCustomMajorTicks') is not None:
        ax_format['xcustom_major_ticks'] = plot_features.get('XMajorTickMarks')

    ax_format['ycustom_major_ticks'] = plot_features.get('YMajorTickMarks')
    ax_format['ycustom_minor_ticks'] = plot_features.get('YMajorTickMarks')
    ax_format['ycustom_minor_ticks'] = plot_features.get('YMajorTickMarks')
    ax_format['plot_xticklabels'] = plot_features.get('PlotXTickLabels')
    ax_format['colormap'] = plot_features.get('ColorMap', 'cividis')

    plot_errorbars = [h.get('ErrorBars', 1) for h in hist_features]
    colors = plt.rcParams['axes.prop_cycle'].by_key()['color']
    cmd.add(f"colors = {str({i: col for i, col in enumerate(colors)})}")

    dataOutPyName = plot_name.strip('/') + '__data.py'
    dataOutPyImport = plot_name.split('/')[-1] + '__data.py'

    cmd.newline()
    cmd.add(f"# the numerical data is stored in a separate file",
            f"dataf = dict()",
            f"exec(open(os.path.split(__file__)[0] + '/{dataOutPyImport}').read(), dataf)")

    ## Fetch numerical data for plotting from separate function
    with open(os.path.join(outdir, dataOutPyName), "a") as dataFile:
        datacmd, data = fd.mkCurves2D(hist_data, True, plot_errorbars, colors)
        cmd.add(datacmd)
        dataFile.write(data)

    if plot_features.get('Legend', 1):
        if plot_features.get('LegendAlign') is None or plot_features.get('LegendAlign') == 'l':
            legend_pos = (float(plot_features.get('LegendXPos', 0.5)),
                          float(plot_features.get('LegendYPos', 0.97)))
            cmd.add(f"legend_pos = {legend_pos}",
                    f"ax.legend(legend_handles, labels, loc='upper left', bbox_to_anchor=legend_pos)")
        if plot_features.get('LegendAlign') == 'r':
            legend_pos = (float(plot_features.get('LegendXPos', 0.97)),
                          float(plot_features.get('LegendYPos', 0.97)))
            cmd.add(f"legend_pos = {legend_pos}""",
                    f"ax.legend(legend_handles, labels, loc='upper right',",
                    f"          bbox_to_anchor=legend_pos,markerfirst=False)")

    # Set text labels on axes
    cmd.newline()
    cmd.add(f"ax.set_xlabel(ax_xLabel)",
            f"ax.set_ylabel(ax_yLabel, loc='top')",
            f"ax.set_title(ax_title, loc='left')",
            f"ax.set_xscale(ax_xScale)",
            f"ax.set_yscale(ax_yScale)")

    cmd.newline()
    cmd.add(f"# tick formatting",
            f"plt.rcParams['xtick.top'] = {plot_features.get('XTwosidedTicks', True)}",
            f"plt.rcParams['ytick.right'] = {plot_features.get('YTwosidedTicks', True)}")

    # https://stackoverflow.com/questions/8213522/when-to-use-cla-clf-or-close-for-clearing-a-plot-in-matplotlib
    # plt.close(fig) instead of fig.clf()?
    for fmt in formats:
        figName = plot_name.split('/')[-1]+f'.{fmt.lower()}'
        cmd.add(f"plt.savefig(os.path.join(plot_outdir, '{figName}'), format='{fmt}')")
    cmd.newline()
    cmd.add(f"plt.close(fig)")
    return cmd.get()


def process(yaml_file, plot_name, outdir, formats):
    """\
    Write an executable python script that only relies on YODA functionalities

    Parameters
    ----------
    yaml_file : TODO
    plot_name : TODO
    outdir    : TODO
    """
    outPyName = os.path.join(outdir, plot_name.strip('/')) + '.py'

    ## Contents to write out to .py script
    mplScript = yoda.util.StringCommand("#! /usr/bin/env python\n")

    yodaVersion = str(yoda.__version__)
    import datetime
    currentTime = datetime.datetime.now().strftime("%d-%m-%Y (%H:%M:%S)")

    ## Add disclaimer and imports
    ## NB: the phrase "Path:" outside a comment, anywhere in a .py file, forces download over inline view!
    mplScript.add(f"# This Python script was auto-generated using YODA v{yodaVersion}.",
                  f"# Analysis object: {plot_name}",
                  f"# Timestamp: {currentTime}")
    mplScript.newline()
    mplScript.add("import matplotlib as mpl",
                  "import matplotlib.pyplot as plt",
                  "mpl.use('Agg') # comment out for interactive use",
                  "import os",
                  "import numpy as np")
    mplScript.newline()

    rivetPlotsDir = os.path.join(os.getcwd(), outdir)
    plotDir = rivetPlotsDir + os.path.split(plot_name)[0]

    mplScript.add("plotDir = os.path.split(os.path.realpath(__file__))[0]")
    mplScript.add("if 'YODA_USER_PLOT_PATH' in globals():",
                  "    plot_outdir = globals()['YODA_USER_PLOT_PATH']",
                  "else:",
                  "    plot_outdir = plotDir")

    hist_data = {}
    has_latex = shutil.which('tex')

    ## Parse yaml file for histogram data
    if isinstance(yaml_file, str): # If the file name of the yaml file is passed...
        hist_data_multiweights = {}
        yaml_dicts = yoda.util._parseyaml(yaml_file)
        hist_data = _parse_yoda_hist(yaml_dicts) # TODO: add multiweights option here too?
    else: # If the dictionary object is passed
        yaml_dicts = yaml_file
        hist_data = {
            k : hist_dict['nominal'] if 'Scatter' in hist_dict['nominal'].type() else \
                hist_dict['nominal'].mkScatter() for k, hist_dict in yaml_dicts['histograms'].items()
        }

        hist_data_band = {}
        hist_data_multiweights = {}
        ## Iterate over input YODA files
        for hist_dict_key, hist_dict in yaml_dicts['histograms'].items():

            ## For each YODA file, store every multiweight scatter object in a dict
            hist_data_multiweights[hist_dict_key] = {k: v for k,v in hist_dict.items() if k.startswith("multiweight")}
            if 'BandUncertainty' in hist_dict.keys():
                hist_data_multiweights[hist_dict_key]['BandUncertainty'] = hist_dict['BandUncertainty']
            elif hist_dict.get('ErrorBands', 0):
                hist_data_multiweights[hist_dict_key]['BandUncertainty'] = hist_dict['nominal'].clone()

    hist_features = [ val for val in yaml_dicts['histograms'].values() ]
    output_filename = os.path.join(outdir, plot_name.strip('/'))
    _preprocess_text(yaml_dicts)

    # convert any Scatter1Ds to Scatter2D
    for k,v in hist_data.items():
        if v.type() == "Scatter1D":
            scatter_2D = yoda.core.Scatter2D()
            scatter_2D.addPoint(1.0, v.point(0).x(), 0.5, v.point(0).xErrs())
            if v.hasAnnotation("IsRef"): scatter_2D.setAnnotation('IsRef', 1)
            hist_data[k] = scatter_2D
            yaml_dicts['plot features']['DummyXaxis'] = True

    if all(h.dim() <= 2 for h in hist_data.values()): # or all(isinstance(h, yoda.Scatter2D) for h in hist_data.values()):
        mplScript.add( mkPlottingScript1D(hist_data, hist_features, yaml_dicts, outdir,
                                          plot_name, hist_data_multiweights, formats, has_latex) )

    elif all(h.dim() == 3 for h in hist_data.values()):
        mplScript.add( mkPlottingScript2D(hist_data.values(), hist_features, yaml_dicts,
                                          outdir, plot_name, formats, has_latex) )
    else:
        print('Error with Class types:', [type(h) for h in hist_data.values()])
        raise NotImplementedError('Class type cannot be plotted yet')

    with open(outPyName, "a") as f:
        f.write(mplScript.get())
    st = os.stat(outPyName)
    os.chmod(outPyName, st.st_mode | stat.S_IEXEC)

    return outPyName


def _parse_yoda_hist(yaml_dicts):
    """Read yoda string and return yoda object."""
    hist_data = []
    for hist_dict in yaml_dicts['histograms'].values():
        with io.StringIO(hist_dict['nominal']) as file_like:
            hist_data.append(yoda.readYODA(
                file_like, asdict=False)[0].mkScatter())
    return hist_data


def _preprocess_text(yaml_dicts):
    """Preprocess text to convert convenient HEP units and other symbols to mathtext."""
    if 'plot features' not in yaml_dicts:
        yaml_dicts['plot features'] = {}
    plot_features = yaml_dicts['plot features']
    for label in plot_features:
        if any(label.endswith(sub) for sub in ('Title', 'XLabel', 'YLabel', 'ZLabel')):
            plot_features[label] = preprocess(plot_features[label])
        if any(label.endswith(sub) for sub in ('XCustomMajorTicks', 'YCustomMajorTicks', 'ZCustomMajorTicks')):
            plot_features[label] = [preprocess(tick) for tick in plot_features[label]]

    for hist_setting in yaml_dicts['histograms'].values():
        if 'Title' in hist_setting:
            hist_setting['Title'] = preprocess(hist_setting['Title'])


class AnyObject(object):
    """Necessary for custom legend handler."""
    pass


class AnyObjectHandler(object):
    """Creates custom legend handler for Data."""

    def legend_artist(self, legend, orig_handle, fontsize, handlebox):
        x0, y0 = handlebox.xdescent, handlebox.ydescent
        width, height = handlebox.width, handlebox.height
        patch = mpl.patches.Circle(
            (x0+width/2, y0+height/2), (2.5)**0.5, facecolor='black')
        handlebox.add_artist(patch)
        patch = mpl.patches.Rectangle(
            (-0.4, 3), width+0.8, 0.8, facecolor='black')
        handlebox.add_artist(patch)
        patch = mpl.patches.Rectangle(
            (width/2-0.4, 0), 0.8, height, facecolor='black')
        handlebox.add_artist(patch)
        return patch
