import os, stat, shutil

import yoda
from yoda.plotting.mlp_preprocessor import preprocess, applyRegexes
import yoda.plotting.fetch_data as fd
import yoda.plotting.utils as putils

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

import math, re
import numpy as np

RE_PATTERN_PATH = re.compile(r'PATH\[(.+?)\]')
RE_PATTERN_NAME = re.compile(r'NAME\[(.+?)\]')


def setCanvas(cmd, fig_dims, margins, nRatios = 0, fig_specs = None):
    """Return Pythonic string corresponding to
    basic mpl canvas, with optional splitting."""
    ## Create fig and axes
    cmd.newline()
    canvasW, canvasH = fig_dims
    cmd.add(f"# Adjust canvas width and height",
            f"canvasW = {canvasW}",
            f"canvasH = {canvasH}",
            f"figW *= canvasW/10.",
            f"figH *= canvasH/9.")
    cmd.newline()
    if nRatios:
        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,")
        fig_specs = tuple(fig_specs)
        cmd.add(f"                  figsize=(figW,figH), gridspec_kw={{'height_ratios': {fig_specs}}})")
    else:
        cmd.add(f"# Create figure and axis objects", f"fig, ax = plt.subplots(1, 1)")
    cmd.newline()
    cmd.add(f"# Set figure margins",
            f"plt.subplots_adjust(",
            f"    left   = {margins[0]} * plt.rcParams['figure.subplot.left'],",
            f"    right  = {margins[1]} * plt.rcParams['figure.subplot.right'],",
            f"    top    = {margins[2]} * plt.rcParams['figure.subplot.top'],",
            f"    bottom = {margins[3]} * plt.rcParams['figure.subplot.bottom'])")
    cmd.newline()


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', {})

    # check if user supplied regex-based label customisation
    re_captures = []
    for item in plot_features.get('PlotRegexes', []):
        for match in RE_PATTERN_NAME.findall(item):
            sub = re.findall(match, list(hist_data.values())[0].name())
            if sub:  re_captures.append(sub[0])
        for match in RE_PATTERN_PATH.findall(item):
            sub = re.findall(match, list(hist_data.values())[0].path())
            if sub:  re_captures.append(sub[0])
    _apply_regexes(plot_features, re_captures, plot_features.get('PlotRegexSubs', {}))

    # adjust font sizes if necessary:
    for key, val in {
        'AxisLabelSize'  : 'axes.labelsize',  # fontsize of the axis labels
        'XTickSize'      : 'xtick.labelsize', # fontsize of the x-tick labels
        'YTickSize'      : 'ytick.labelsize', # fontsize of the y-tick labels
        'TitleSize'      : 'axes.titlesize',  # fontsize of the figure title
        'LegendFontSize' : 'legend.fontsize', # fontsize of the legend labels
      }.items():
        if key not in plot_features:  continue
        cmd.add(f"plt.rcParams['{val}'] = {plot_features[key]}")

    if not has_latex:
        cmd.newline()
        cmd.add(f"plt.rcParams['text.usetex'] = False",
                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", "")}'""")

    ax_format = {}  # Stores the items for formatting the axes in a dict
    ax_format['logx'] = int(plot_features.get('LogX', 0))
    ax_format['logy'] = int(plot_features.get('LogY', 1))

    hasNonZero = any([np.any(h.yVals()) for h in hist_data.values()])
    xscale = 'log' if ax_format['logx'] else 'linear'
    yscale = 'log' if ax_format['logy'] and hasNonZero else 'linear'
    cmd.add(f"""ax_xScale = '{xscale}'""",
            f"""ax_yScale = '{yscale}'""")

    ## Work out xMin and xMax from all curves, excluding under/overflow infinities
    nOverflows = [ h.annotation('nOverflows', 0) for h in hist_data.values() ]
    # keeping overflows requires user to have set IncludeOverflows=1 and all objects having two overflows
    hasOverflows = set(nOverflows).pop() == 2 if int(plot_features.get('IncludeOverflows', '0')) else False
    xmin = plot_features.get('XMin', min([min(h.xMins()[1:-1] \
                                          if nOverflows[i] == 2 and int(plot_features.get('IncludeOverflows', '0')) else \
                                          h.xMins()) for i,h in enumerate(hist_data.values())]))
    xmax = plot_features.get('XMax', max([max(h.xMaxs()[1:-1] \
                                          if nOverflows[i] == 2 and int(plot_features.get('IncludeOverflows', '0')) else \
                                          h.xMaxs()) for i,h in enumerate(hist_data.values())]))
    if ax_format['logx']:
        xmin = np.log10(xmin)
        xmax = np.log10(xmax)

    # work out overflow bin width as fraction of overall binned range
    overflowWidth = float(plot_features.get('OverflowWidthFrac', 0.07)) * (xmax - xmin) if hasOverflows else 0.0
    underflowWidth = float(plot_features.get('UnderflowWidthFrac', 0.07)) * (xmax - xmin) if hasOverflows else 0.0

    # adjust limits
    xmin -= underflowWidth
    xmax += overflowWidth

    ## Set plot limits
    ax_format['xlim'] = (float(10 ** xmin if ax_format['logx'] else xmin),
                         float(10 ** xmax if ax_format['logx'] else xmax))
    if any(nOverflows):
        # set dummy xMin/xMax in case overflows are to be plotted
        putils.adjustOverflows(hist_data.values(), ax_format['xlim'],
                               int(plot_features.get('IncludeOverflows', '0')),
                               hasOverflows and underflowWidth > 0,
                               hasOverflows and overflowWidth > 0)
        for item in multiweight_dict.values():
            putils.adjustOverflows(item.values(), ax_format['xlim'],
                                   int(plot_features.get('IncludeOverflows', '0')),
                                   hasOverflows and underflowWidth > 0,
                                   hasOverflows and overflowWidth > 0)

    ## 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 = float(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 if max_ymax > 0.0 else 0.9*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 = float(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 0.11*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)

    if math.isclose(ymin, ymax):
        ymax = 10*ymin if ymin else 1.0

    ax_format['ylim'] = (float(ymin), float(ymax))

    leg_pass = False
    leg_titles, leg_curves, leg_pos, ratio_panels = {}, {}, {}, {}
    plot_features.setdefault('Legend', 1)
    for label in plot_features.keys():
        # Check for legend configs
        if re.fullmatch(r"Legend(\d+)?$", label) and plot_features.get(label, 1):
            leg_curves[label] = plot_features.get(f'{label}Only', '').split()
            leg_titles[label] = (list(), plot_features.get(f'{label}Title', ''))
            for i, name in enumerate(hist_data.keys()):
                if leg_curves[label] and not name in leg_curves[label]:  continue
                lentry = str(hist_features[i].get('Title', 'Curve {}'.format(i+1)))
                lentry = lentry.replace('.yoda.gz','').replace('.yoda', '')
                if len(leg_titles[label][0]) and lentry == '':  continue
                elif lentry == '':  lentry = 'Curve 0'
                if lentry.count('_') and lentry.count('$') < 2:
                    lentry = lentry.replace('_', r'\_')
                if plot_features.get('GofLegend', False) and i > 0:
                    vals1, vals2, errs1, errs2 = putils.selectValRange(next(iter(hist_data.values())),
                                                                       hist_data[name],
                                                                       *ax_format['xlim'])
                    gof = None; gofLabel = ''
                    gofType = plot_features.get('GofType', 'naiveChi2reduced')
                    if gofType == 'naiveChi2reduced':
                        gof = yoda.naiveChi2reduced(vals1, vals2, errs1, errs2)
                        gofLabel = r'\chi^2/N_\mathrm{df}'
                    elif gofType == 'naiveChi2':
                        gof = yoda.naiveChi2(vals1, vals2, errs1, errs2)
                        gofLabel = r'\chi^2'
                    else:
                        raise NotImplementedError(f'Gof type "{gofType}" not implemented yet')
                    lentry += f', ${gofLabel}={gof:.2f}$'
                leg_pass = True
                leg_titles[label][0].append(lentry)
                if label not in leg_pos:
                    posx, posy, anc, al = putils.legendDefaults(list(hist_data.values())[0],
                                                                ax_format['xlim'], ax_format['ylim'],
                                                                ax_format['logx'], ax_format['logy'])
                    leg_pos[label] = ((float(plot_features.get(f'{label}XPos', posx)),
                                       float(plot_features.get(f'{label}YPos', posy))),
                                       plot_features.get(f'{label}Anchor', anc),
                                       plot_features.get(f'{label}Align', al))

        # Check for ratio-panel configs
        if not re.fullmatch(r"RatioPlot(\d+)?$", label):  continue
        if not plot_features.get(label, 0):  continue
        default = label == 'RatioPlot' and plot_features.get('RatioPlot', 1)
        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():
                if label in ratio_panels:
                    ratio_panels[label].append( rcurve )
                else:
                    ratio_panels[label] = [ 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)
    if not ratio_panels and plot_features.get('RatioPlot', 1):
        ## Unless user specifically asks not have ratio panel,
        ## have at least one by default
        ratio_panels['RatioPlot'] = list(hist_data.keys())
        ax_format['ratio0_logy'] = plot_features.get('RatioPlotLogY', 0)

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

    # Check if user set custom x-tick marks
    if plot_features.get('XCustomMajorTicks') is not None:
        # clean tabs
        #tmp = re.split(r'\s+', ''.join(plot_features.get('XCustomMajorTicks')))
        tmp = plot_features.get('XCustomMajorTicks').split('\t')
        # create map of tick positon <-> tick label
        tmp = { float(t) : l for t, l in zip(*[iter(tmp)]*2) }
        cmd.add("xMajorTickMarks = {}".format(list(tmp.keys())),
                "xCustomTickLabels = {}".format(list(tmp.values())))
    elif plot_features.get('XMajorTickMarks') is not None:
        tmp = [ float(x) for x in plot_features.get('XMajorTickMarks').split() ]
        cmd.add("xMajorTickMarks = {}".format(tmp))
    if plot_features.get('XMinorTickMarks', ''):
        cmd.add(f"xMinorTickMarks = { int(plot_features.get('XMinorTickMarks')) }")

    # Check if user set custom y-tick marks
    if plot_features.get('YCustomMajorTicks') is not None:
        # clean tabs
        tmp = [ y for y in plot_features.get('YCustomMajorTicks') if y != '\t' ]
        # create map of tick positon <-> tick label
        tmp = { float(t) : l for t, l in zip(*[iter(tmp)]*2) }
        cmd.add("yMajorTickMarks = {}".format(list(tmp.keys())),
                "yCustomTickLabels = {}".format(list(tmp.values())))
    elif plot_features.get('YMajorTickMarks') is not None:
        tmp = [ float(y) for y in plot_features.get('YMajorTickMarks').split() ]
        cmd.add("yMajorTickMarks = {}".format(tmp))
    if plot_features.get('YMinorTickMarks') is not None:
        cmd.add(f"yMinorTickMarks = { int(plot_features.get('YMajorTickMarks')) }")

    # Add a dictionary with TeX-friendly legend labels
    if leg_pass:
        cmd.newline()
        cmd.add(f"# TeX-friendly labels for the legend", "labels = {")
        for legName, (labels, title) in leg_titles.items():
            texLabels = 'r"' + '", r"'.join(labels) + '"' #< Use r-strings to be TeX friendly
            texTitle = 'r"' + title + '"'
            cmd.add(f"    '{legName.lower()}' : ([ {texLabels} ], {texTitle}),")
        cmd.add("}")


    # determine figure specs
    fig_specs = [ float(plot_features.get('PlotSizeY', 6)) ]
    for label in ratio_panels:
        fig_specs.append(float(plot_features.get(label+'YSize', 3)))
    fig_dims = (float(plot_features.get('PlotSizeX', 10)), sum(fig_specs))
    margins = (plot_features.get('LeftMargin',   1.0),
               plot_features.get('RightMargin',  1.0),
               plot_features.get('TopMargin',    1.0),
               plot_features.get('BottomMargin', 1.0))

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

    fc = plot_features.get('FrameColor', None)
    if fc:
        cmd.add(f"fig.set_facecolor('{fc}')\n")

    # default ticks rotations
    xRot = plot_features.get('XMajorTicksAngle', 0)
    yRot = plot_features.get('YMajorTicksAngle', 0)

    ## Set some ratio-panel formatting defaults
    deviation_modes = []
    for i, key in enumerate(ratio_panels.keys()):
        ratioName = 'ratio%d_ax' % i
        ## Ratio plot default y-ranges
        ratioYMin = float(plot_features.get(key+'YMin', 0.5))
        ratioYMax = float(plot_features.get(key+'YMax', 1.4999))  # avoid 1.5 tick mark
        rpmode = plot_features.get(key+'Mode', 'mcdata')
        deviation_modes.append(rpmode == 'deviation')
        if deviation_modes[-1]:
            cmd.add(f"{ratioName}.set_yticks(list(range(-3,3)),",
                    f"                     [r'$-3\\,\\sigma$', r'$-2\\,\\sigma$', r'$-1\\,\\sigma$',",
                    f"                      r'$0\\,\\sigma$', r'$1\\,\\sigma$', r'$2\\,\\sigma$'])",
                    f"{ratioName}.set_ylim(-3.0, 3.0)")
        else:
            ratioTickDist = (ratioYMax-ratioYMin)/5.0 # default: five tick labels, 1 significant digit:
            ratioTickDist = round(ratioTickDist, -int(math.floor(math.log10(abs(ratioTickDist)))))
            cmd.add(f"{ratioName}.yaxis.set_major_locator(mpl.ticker.MultipleLocator({ratioTickDist}))",
                    f"{ratioName}.set_ylim({ratioYMin}, {ratioYMax})")
        ratioYLabel = plot_features.get(key+'YLabel', 'Ratio')
        cmd.add(f"{ratioName}.set_ylabel('{ratioYLabel}')")
        if plot_features.get('DummyXaxis'):
            cmd.add(f"{ratioName}.tick_params(axis='x', which='both', bottom=False, top=False, labelbottom=False)")

        if plot_features.get('XCustomMajorTicks') is not None:
            cmd.add(f"{ratioName}.set_xticks(xMajorTickMarks, xCustomTickLabels, rotation={xRot}, minor=False)",
                    f"{ratioName}.tick_params(axis='x', which='minor', bottom=False, top=False)")
        elif plot_features.get('XMajorTickMarks') is not None:
            cmd.add(f"{ratioName}.set_xticks(xMajorTickMarks, rotation={xRot}, minor=False)")
            if not plot_features.get('XMinorTickMarks', ''):
                cmd.add(f"{ratioName}.tick_params(axis='x', which='minor', bottom=False, top=False)")
        if plot_features.get('XMinorTickMarks', ''):
            n = int(plot_features.get('XMinorTickMarks'))
            if n == 0:
                cmd.add(f"{ratioName}.xaxis.set_minor_locator(mpl.ticker.NullLocator())")
            else:
                cmd.add(f"{ratioName}.xaxis.set_minor_locator(mpl.ticker.MultipleLocator({n}))")

    plot_errorbars = [h.get('ErrorBars', 1) 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 pre in [''] + list(ratio_panels.keys()):
            for base_tag in ['Title', 'LineStyle', 'LineWidth', 'MarkerStyle', 'MarkerSize',
                             'ConnectBins', 'ConnectMarkers', 'FillColor', 'FillOpacity',
                             'ErrorCapSize', 'ErrorBand', 'ErrorBandColor',
                             'ErrorBandStyle', 'ErrorBandOpacity' ]:
                tag = pre+base_tag
                if tag in hist_features[i]:
                    ao.setAnnotation(tag, hist_features[i][tag])
                    del hist_features[i][tag]
                elif tag == 'ConnectBins' and ao.annotation('IsRef', 0):
                    ao.setAnnotation(tag, 0)
        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()

    ## 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)")

    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_modes, **ax_format)
        cmd.add(datacmd)
        dataOutFile.write(data)

    ## Fill, position and draw legend
    for legName, legCurves in leg_curves.items():
        if legCurves:
            cmd.add(f"legend_items = [")
            for lc in legCurves:
                cmd.add(f"    legend_handles['{lc}'],")
            cmd.add(f"]")
        else:
            cmd.add(f"legend_items = list(legend_handles.values())")
        legXY, legAnchor, legAlign = leg_pos[legName]
        markerPos = ',markerfirst=False' if legAlign == 'r' else ''
        titleAlign = 'right' if legAlign == 'r' else 'left'
        cmd.add(f"ax.add_artist(ax.legend(legend_items,",
                f"                        labels['{legName.lower()}'][0],",
                f"                        title=labels['{legName.lower()}'][1],",
                f"                        alignment='{titleAlign}',",
                f"                        loc='{legAnchor}',",
                f"                        bbox_to_anchor={legXY}{markerPos}))")

    ## Set text labels on axes
    cmd.newline()
    cmd.add("# set plot metadata as defined above")
    xLabelSep = float(plot_features.get(f'XLabelSep', -0.15 if len(ratio_panels.keys()) else -0.05 ))
    if len(ratio_panels.keys()):
        # only put the label for the last ratio panel
        ratioName = 'ratio%d_ax' % (len(ratio_panels.keys()) - 1)
        cmd.add(f"{ratioName}.set_xlabel(ax_xLabel)",
                f"{ratioName}.xaxis.set_label_coords(1., {xLabelSep})")
    else:
        # no ratio panels, so put on main panel
        cmd.add(f"ax.set_xlabel(ax_xLabel)",
                f"ax.xaxis.set_label_coords(1., {xLabelSep})")

    yLabelSep = float(plot_features.get(f'YLabelSep', -0.11))
    cmd.add(f"ax.set_ylabel(ax_yLabel, ha='right', va='top')",
            f"ax.yaxis.set_label_coords({yLabelSep}, 1.0)",
            f"ax.set_title(ax_title, loc='left')",
            f"ax.set_xscale(ax_xScale)",
            f"ax.set_yscale(ax_yScale)")

    ## Toggle x/y lims
    cmd.add("ax.set_xlim(xLims)")
    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)}")

    ## Set custom x-tick marks if applicable
    if plot_features.get('XCustomMajorTicks') is not None:
        cmd.add(f"ax.set_xticks(xMajorTickMarks, xCustomTickLabels, rotation={xRot}, minor=False)",
                f"ax.tick_params(axis='x', which='minor', bottom=False, top=False)")
    elif plot_features.get('XMajorTickMarks') is not None:
        cmd.add(f"ax.set_xticks(xMajorTickMarks, rotation={xRot}, minor=False)")
        if not plot_features.get('XMinorTickMarks', ''):
            cmd.add(f"ax.tick_params(axis='x', which='minor', bottom=False, top=False)")
    if plot_features.get('XMinorTickMarks', ''):
        n = int(plot_features.get('XMinorTickMarks'))
        if n == 0:
            cmd.add(f"ax.xaxis.set_minor_locator(mpl.ticker.NullLocator())")
        else:
            cmd.add(f"ax.xaxis.set_minor_locator(mpl.ticker.MultipleLocator({n}))")

    ## Set custom y-tick marks if applicable
    if plot_features.get('YCustomMajorTicks') is not None:
        cmd.add(f"ax.set_yticks(yMajorTickMarks, yCustomTickLabels, rotation={yRot}, minor=False)",
                f"ax.tick_params(axis='y', which='minor', bottom=False, top=False)")
    elif plot_features.get('YMajorTickMarks') is not None:
        cmd.add(f"ax.set_yticks(yMajorTickMarks, rotation={yRot}, minor=False)")
        if plot_features.get('YMinorTickMarks') is None:
            cmd.add(f"ax.tick_params(axis='y', which='minor', bottom=False, top=False)")
    if plot_features.get('YMinorTickMarks') is not None:
        n = int(plot_features.get('YMinorTickMarks'))
        if n == 0:
            cmd.add(f"ax.yaxis.set_minor_locator(mpl.ticker.NullLocator())")
        else:
            cmd.add(f"ax.yaxis.set_minor_locator(mpl.ticker.MultipleLocator({n}))")

    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))")

    # Check if user disabled tick labels
    xyticklabels = plot_features.get('PlotTickLabels',  1)
    xticklabels  = plot_features.get('PlotXTickLabels', 1)
    yticklabels  = plot_features.get('PlotYTickLabels', 1)
    if not xyticklabels or not xticklabels:
       cmd.add("ax.xaxis.set_tick_params(labelbottom=False)")
    if not xyticklabels or not yticklabels:
       cmd.add("ax.yaxis.set_tick_params(labelleft=False)")

    # Annotate under-/overflow bins with labels
    if underflowWidth > 0. or overflowWidth > 0.:
        cmd.add(f"\n# Annotate under-/overflow bins")
        ref_label = list(hist_data.keys())[0]
        ylims = np.log10(ax_format['ylim']) if ax_format['logy'] else ax_format['ylim']
        underflow_ypos = None
        overflow_ypos = None
        for h in hist_data.values():
            y = h.yVals()[0]
            if math.isfinite(y) and (y > 0. or not ax_format['logy']):
                underflow_ypos = np.log10(y) if ax_format['logy'] else y
            y = h.yVals()[-1]
            if math.isfinite(y) and (y > 0. or not ax_format['logy']):
                overflow_ypos = np.log10(y) if ax_format['logy'] else y
            if underflow_ypos is not None and overflow_ypos is not None:
                break
        if underflowWidth > 0.:
            if underflow_ypos is None:
                underflow_ypos = 0.5*sum(ylims)
            flip = int(underflow_ypos < 0.75*ylims[0] + 0.25*ylims[1])
            sf = float(plot_features.get('UnderflowLabelOffset', 0.03))
            underflow_ypos += sf*(ylims[1] - ylims[0]) * (-1.0) ** flip # shift above/below bin value
            if ax_format['logy']:  underflow_ypos = 10 ** underflow_ypos
            underflow_label = plot_features.get('UnderflowLabel', r'$-\infty\leq$')
            underflow_color = plot_features.get('UnderflowLabelColor', 'black')
            underflow_size = plot_features.get('UnderflowLabelSize', 8)
            cmd.add(f"underflow_xpos = 0.5*sum(dataf['xedges']['{ref_label}'][:2])",
                    f"underflow_ypos = {float(underflow_ypos):.2e}",
                    f"ax.annotate(r'{underflow_label}', xy=(underflow_xpos, underflow_ypos),",
                    f"            ha='center', fontsize={underflow_size}, color='{underflow_color}')")
        if overflowWidth > 0.:
            if overflow_ypos is None:
                overflow_ypos = 0.5*sum(ylims)
            flip = int(overflow_ypos < 0.75*ylims[0] + 0.25*ylims[1])
            sf = float(plot_features.get('OverflowLabelOffset', 0.03))
            overflow_ypos += sf*(ylims[1] - ylims[0]) * (-1.0) ** flip # shift above/below bin value
            if ax_format['logy']:  overflow_ypos = 10 ** overflow_ypos
            overflow_label = plot_features.get('OverflowLabel', r'$\leq\infty$')
            overflow_color = plot_features.get('OverflowLabelColor', 'black')
            overflow_size = plot_features.get('OverflowLabelSize', 8)
            cmd.add(f"overflow_xpos = 0.5*sum(dataf['xedges']['{ref_label}'][-2:])",
                    f"overflow_ypos = {float(overflow_ypos):.2e}",
                    f"ax.annotate(r'{overflow_label}', xy=(overflow_xpos, overflow_ypos),",
                    f"            ha='center', fontsize={overflow_size}, color='{overflow_color}')")

    for key in ratio_panels.keys():
        ratioName = 'ratio' + ('0' if 't' == key[-1] else key[-1])
        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))")
        ratioXYticklabels = plot_features.get(key+'TickLabels',  1)
        ratioXticklabels  = plot_features.get(key+'XTickLabels', 1)
        ratioYticklabels  = plot_features.get(key+'YTickLabels', 1)
        if not ratioXYticklabels or not ratioXticklabels:
            cmd.add(f"{ratioName}_ax.xaxis.set_tick_params(labelleft=False)")
        if not ratioXYticklabels or not ratioYticklabels:
            cmd.add(f"{ratioName}_ax.yaxis.set_tick_params(labelleft=False)")
    if plot_features.get('DummyXaxis'):
        cmd.add("ax.tick_params(axis='x', which='both', bottom=False, top=False, labelbottom=False)")

    # Loop over text boxes
    for tb in [ k for k in plot_features if k.startswith('text_') ]:
        textBoxName = tb.replace('text_', '')
        textBox = plot_features[tb]
        # Loop over panels
        for i, key in enumerate(list(ratio_panels.keys()) + ['Plot']):
            panelName = 'ax' if key == 'Plot' else 'ratio%d_ax' % i
            if not textBox.get(key+'Text', ''):  continue
            txt = textBox.get(key+'Text', '')
            xpos = textBox.get('XPos', 0.5)
            ypos = textBox.get('YPos', 0.5)
            angle = textBox.get('Rotation', None)
            fontsize = textBox.get('FontSize', None)
            boxopts = f', transform={panelName}.transAxes'
            if fontsize:
                boxopts += f', fontsize={fontsize}'
            if angle:
                boxopts += f', rotation={angle}, rotation_mode="anchor"'
            # TODO: add more formatting options?
            cmd.add(f"{panelName}.text({xpos}, {ypos}, '{txt}'{boxopts})")

    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', {})

    # check if user supplied regex-based label customisation
    re_captures = []
    for item in plot_features.get('PlotRegexes', []):
        for match in RE_PATTERN_NAME.findall(item):
            sub = re.findall(match, list(hist_data.values())[0].name())
            if sub:  re_captures.append(sub[0])
        for match in RE_PATTERN_PATH.findall(item):
            sub = re.findall(match, list(hist_data.values())[0].path())
            if sub:  re_captures.append(sub[0])
    _apply_regexes(plot_features, re_captures, plot_features.get('PlotRegexSubs', {}))

    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
    xmin = plot_features.get('XMin', min([min(h.xMins()) for h in hist_data.values()]))
    xmax = plot_features.get('XMax', max([max(h.xMaxs()) for h in hist_data.values()]))
    ymin = plot_features.get('YMin', min([min(h.yMins()) for h in hist_data.values()]))
    ymax = plot_features.get('YMax', max([max(h.yMaxs()) for h in hist_data.values()]))
    zmin = plot_features.get('ZMin', min([min(h.zMins()) for h in hist_data.values()]))
    zmax = plot_features.get('ZMax', max([max(h.zMaxs()) for h in hist_data.values()]))
    ax_format['xlim'] = (float(xmin), float(xmax))
    ax_format['ylim'] = (float(ymin), float(ymax))
    ax_format['zlim'] = (float(zmin), float(zmax))

    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))

    cmd.add(f"xLims = {ax_format['xlim']}")
    cmd.add(f"yLims = {ax_format['ylim']}")
    cmd.add(f"zLims = {ax_format['zlim']}")

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

    fig_dims = (plot_features.get('PlotSizeX', 10),
                plot_features.get('PlotSizeY', 10))
    margins = (plot_features.get('LeftMargin',   1.0),
               plot_features.get('RightMargin',  1.0),
               plot_features.get('TopMargin',    1.0),
               plot_features.get('BottomMargin', 1.0))
    setCanvas(cmd, fig_dims, margins)

    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('XCustomMajorTickMarks')

    ax_format['ycustom_major_ticks'] = plot_features.get('YMajorTickMarks')
    ax_format['ycustom_minor_ticks'] = plot_features.get('YMinorTickMarks')
    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.newline()
    cmd.add("# Color map: curve index -> color",
            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(2)
    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, **ax_format)
        cmd.add(datacmd)
        dataFile.write(data)

    if plot_features.get('Legend', 0):
        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)")

    ## Toggle x/y lims
    cmd.add("ax.set_xlim(xLims)")
    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)}")

    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 = ['pdf']):
    """\
    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 python3\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",
                  "from numpy import nan, inf")
    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)
        yaml_dicts.setdefault('plot features', {})
        hist_data = _parse_yoda_hist(yaml_dicts) # TODO: add multiweights option here too?
    else: # If the dictionary object is passed
        yaml_dicts = yaml_file
        yaml_dicts.setdefault('plot features', {})
        for k, hist_dict in yaml_dicts['histograms'].items():
            hasOverflows = int(yaml_dicts['plot features'].get('IncludeOverflows', '0'))
            hist_data[k] = putils.mkPlotFriendlyScatter(hist_dict['nominal'], includeOverflows=hasOverflows)
            if 'Scale' in hist_dict:
                hist_data[k].scale(hist_data[k].dim()-1, float(hist_dict['Scale']))
            if 'Rebin' in yaml_dicts['plot features']:
                try:
                    newao = putils.scatter_rebin(hist_data[k], yaml_dicts['plot features'].get('Rebin'))
                    hist_data[k] = newao
                except:
                    print('WARNING - rebinning attempt failed!')
            if hist_data[k].hasAnnotation('DummyXaxis'):
                yaml_dicts['plot features']['DummyXaxis'] = True
            for A in 'XYZ':
                anno = A + 'CustomMajorTicks'
                if hist_data[k].hasAnnotation(anno) and not anno in yaml_dicts['plot features']:
                    yaml_dicts['plot features'][anno] = hist_data[k].annotation(anno)

        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:
                hist_data_multiweights[hist_dict_key]['BandUncertainty'] = hist_dict['BandUncertainty']
            for k,v in hist_data_multiweights[hist_dict_key].items():
                if 'Scale' in hist_dict:
                    v.scale(v.dim()-1, float(hist_dict['Scale']))
                try:
                    hist_data_multiweights[hist_dict_key][k] = putils.scatter_rebin(v, yaml_dicts['plot features'].get('Rebin'))
                except:
                    pass

    if not os.path.exists(outdir):  os.makedirs(outdir)
    hist_features = list(yaml_dicts['histograms'].values())
    output_filename = os.path.join(outdir, plot_name.strip('/'))
    _preprocess_text(yaml_dicts)

    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, hist_features, yaml_dicts,
                                          outdir, plot_name, formats, has_latex) )
    else:
        print('Error with Class types:', [h.type() for h in hist_data.values()])
        raise NotImplementedError('Class type or mix of class types 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:
            s = putils.mkPlotFriendlyScatter(yoda.readYODA(file_like, asdict=False)[0])
            hist_data.append(s)
            for A in 'XYZ':
                anno = A + 'CustomMajorTicks'
                if s.hasAnnotation(anno):
                    yaml_dicts['plot features'][anno] = s.annotation(anno)

    return hist_data


def _preprocess_text(yaml_dicts):
    """Preprocess text to convert convenient HEP units and other symbols to mathtext."""
    plot_features = yaml_dicts.get('plot features', {})
    for label in plot_features:
        if any(label.endswith(sub) for sub in ('Title', 'Label', 'CustomMajorTicks')):
            plot_features[label] = preprocess(plot_features[label])

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


def _apply_regexes(features, regex_captures, regex_subs):
    """Apply regular expressions to plot labels, titles and such using the applyRegex() helper method."""
    for label in features:
        if any(label.endswith(sub) for sub in ('Title', 'Label', 'CustomMajorTicks')):
            features[label] = applyRegexes(features[label], regex_captures, regex_subs)


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
