import matplotlib.pyplot as plt
from matplotlib.pyplot import Axes
import numpy as np
from collections.abc import Iterable

import yoda
from yoda import Scatter1D, Scatter2D, Scatter3D, Histo1D, Histo2D

def mkCurves2D(aos_dict, plot_ref=True, error_bars=True, colors=[],
              line_styles=['-', '--', '-.', ':'], contourlevel=None, **kwargs):
    """
    Auxiliary method that constructs and returns two text strings: one which represents
    the numerical data associated with the curves as well as edges, and one which
    contains the logic to plot these for a 2D plot.
    """

    data = yoda.util.StringCommand()
    cmd = yoda.util.StringCommand()

    try:
        aos = list(aos_dict.values())
    except:
        aos = list(aos_dict)
    if not isinstance(aos_dict, Iterable):  # Convert single hist object to list
        aos = [aos_dict]

    if isinstance(error_bars, bool):  # Convert to list of bool vals
        error_bars = [error_bars] * len(aos)

    if not colors:  # Use default mpl prop cycle vals
        colors = plt.rcParams['axes.prop_cycle'].by_key()['color']

    line_properties = LineProperties(colors, line_styles)

    # Define useful variables
    xpoints   = 0.5*(aos[0].xMins() + aos[0].xMaxs())
    xedges    = np.append(aos[0].xMins(), aos[0].xMax())
    ref_xvals = aos[0].xVals()
    ref_xerrs = aos[0].xErrs()
    ypoints   = 0.5*(aos[0].yMins() + aos[0].yMaxs())
    yedges    = np.append(aos[0].yMins(), aos[0].yMax())
    ref_yvals = aos[0].yVals()
    ref_yerrs = aos[0].yErrs()
    zpoints   = 0.5*(aos[0].zMins() + aos[0].zMaxs())
    zedges    = np.append(aos[0].zMins(), aos[0].zMax())
    ref_zvals = aos[0].zVals()
    ref_zerrs = aos[0].zErrs()
    data.add(f"xpoints = {[val for val in xpoints]}",
             f"xedges  = {[val for val in xedges]}",
             f"xmins   = {[val for val in aos[0].xMins()]}",
             f"xmaxs   = {[val for val in aos[0].xMaxs()]}",
             f"xerrs = [",
             f"  [abs(xpoints[i] - xmins[i]) for i in range(len(xpoints))],",
             f"  [abs(xmaxs[i] - xpoints[i]) for i in range(len(xpoints))]",
             f"]")
    data.newline()

    data.add(f"ref_xvals     = {[val for val in ref_xvals]}",
             f"ref_xerrminus = {[err[0] for err in ref_xerrs]}",
             f"ref_xerrplus  = {[err[1] for err in ref_xerrs]}",
             f"ref_xerrs     = [ref_xerrminus, ref_xerrplus]")
    data.newline()

    data.add(f"ypoints   = {[val for val in ypoints]}",
             f"yedges    = {[val for val in yedges]}",
             f"ymins     = {[val for val in aos[0].yMins()]}",
             f"ymaxs     = {[val for val in aos[0].yMaxs()]}",
             f"yerrs = [",
             f"  [abs(ypoints[i] - ymins[i]) for i in range(len(ypoints))],",
             f"  [abs(ymaxs[i] - ypoints[i]) for i in range(len(ypoints))]",
             f"]")
    data.newline()

    data.add(f"ref_yvals     = {[val for val in ref_yvals]}",
             f"ref_yerrminus = {[err[0] for err in ref_yerrs]}",
             f"ref_yerrplus  = {[err[1] for err in ref_yerrs]}",
             f"ref_yerrs     = [ref_yerrminus, ref_yerrplus]")
    data.newline()

    data.add(f"zpoints   = {[val for val in zpoints]}",
             f"zedges    = {[val for val in zedges]}",
             f"zmins     = {[val for val in aos[0].zMins()]}",
             f"zmaxs     = {[val for val in aos[0].zMaxs()]}")
    data.newline()

    data.add(f"ref_zvals     = {[val for val in ref_zvals]}",
             f"ref_zerrminus = {[err[0] for err in ref_zerrs]}",
             f"ref_zerrplus  = {[err[1] for err in ref_zerrs]}",
             f"ref_zerrs     = [ref_zerrminus, ref_zerrplus]")
    data.newline()


    # Plot the reference histogram data
    cmd.add("legend_handles = [] # keep track of handles for the legend")
    # if contourlevel is None:
    #     plt_cntr = False
    # elif not isinstance(contourlevel, list):
    #     plt_cntr = True
    #     contourlevel = list(contourlevel)
    # else:
    #     plt_cntr = True
    colormap = kwargs['colormap'] if 'colormap' in kwargs else 'cividis'

    if plot_ref:
        #suff = ",\n                    norm=mpl.colors.LogNorm())" \
        #suff = ",\n                    locator=mpl.ticker.LogLocator())" \
        suff = ",\n                   norm=mpl.colors.LogNorm(vmin=Z.min(), vmax=Z.max()))" \
               if kwargs['logz'] else ")"
        cmd.add(f"# reference data in main panel",
                f"cmap = '{colormap}'", # TODO This needs to be configurable from the outside
                f"xbin_cent = np.unique(dataf['xpoints'])",
                f"ybin_cent = np.unique(dataf['ypoints'])",
                f"X, Y = np.meshgrid(xbin_cent, ybin_cent)",
                f"Z = np.array(dataf['zpoints']).reshape(X.shape[::-1])",
                f"pc = ax.pcolormesh(X, Y, Z.T, cmap=cmap, shading='auto'"+suff,
                f"cbar = fig.colorbar(pc, orientation='vertical', ax=ax, pad=0.01)",#+suff,
                f"cbar.set_label(ax_zLabel)")

    #
    c_zvals, c_zerrs, c_zups, c_zdowns, c_styles = {}, {}, {}, {}, {}
    if contourlevel is not None:
        cmd.add(f"base_contour = ax.contour(X, Y, zvals.T, levels={contourlevel},",
                f"                          linestyles='--', alpha=0.6, linewidths=2,",
                f"                          colors='black', zorder=2)",
                f"legend_handles += [base_contour]")

        for i, ao in enumerate(aos):
            if i == 0:
                continue
            colidx, linestyle = next(line_properties)
            zorder = 5 + i
            c_zvals = ao.zVals()
            # TODO Why is this added twice?
            c_zvals[ f'c{i}'] = c_zvals
            c_zerrs[ f'c{i}'] = c_zvals
            c_zups[  f'c{i}'] = [err[1] for err in ao.zErrs()] if error_bars[i] else [0] * len(c_zvals)
            c_zdowns[f'c{i}'] = [err[0] for err in ao.zErrs()] if error_bars[i] else [0] * len(c_zvals)
            c_styles[f'c{i}'] = {'color' : f'colors[{colidx}]', 'linestyle' : linestyle, 'zorder' : zorder}

    # write dictionary of labels + yvalues to iterate over in executable python script
    # TODO What is this meant to be? Another z-error? We already added them?
    data.add("zvals = {")
    for label, zhists in c_zvals.items():
        data.add(f"  '{label}' : {[val for val in zvals]},")
    data.add("}", "zerrs = {")
    for label, zrrs in c_zerrs.items():
        data.add(f"  '{label}' : {[val for val in zerrs]},")
    data.add("}", "zups = {")
    for label, zups in c_zups.items():
        data.add(f"  '{label}' : {[val for val in zups]},")
    data.add("}", "z_down = {")
    for label, zdowns in c_zdowns.items():
        data.add(f"  '{label}' : {[val for val in zdowns]},")
    data.add("}")

    cmd.add("# style options for curves",
            "# starts at zorder>=5 to draw curve on top of legend")
    if c_styles:
        cmd.add("styles = {")
        for i, (key, val) in enumerate(c_styles.items()):
            cmd.add(f"  '{key}': " + "{")
            for ky, vl in val.items():
                if ky == 'color':
                    cmd.add(f"""  '{ky}' : {vl.strip("'")}, """)
                elif isinstance(vl, str):
                    cmd.add(f"""  '{ky}' : '{vl}', """)
                else:
                    cmd.add(f"""  '{ky}' : {vl}, """)
            if i < len(c_styles) - 1:
                cmd.add("   },")
            else:
                cmd.add("}}")

        cmd.add("# curve from input yoda files in main panel",
                f"for label in dataf['zerrs'].keys():",
                f"    temp_zvals = np.array(dataf['c_zvals'][label]).reshape(X.shape[::-1])",
                f"    tmp_contour = ax.contour(X, Y, temp_zvals,",
                f"                             levels={contourlevel},",
                f"                             linestyle=c_styles[label]['linestyle'],",
                f"                             linewidths=2, colors=c_styles[label]['color'],",
                f"                              zorder=c_styles[label]['zorder'])",
                f"    legend_handles += [tmp_contour]")

    cmd.newline()

    return cmd.get(), data.get()


def nanpad(xprobe, xref, yvals, yerrs = None):
    """
    match the length of x-values, y-values and errors between two different analysis objects
    if one is longer than the other.
    If one list has missing x-values, the x-values will be added in with nan for y-values and errors.
    """

    padded_yvals = np.array([ yvals[x==xprobe][0] if x in xprobe else \
                              np.nan for x in xref ])
    padded_yerrs = None if yerrs is None else \
                   np.array([ [errs[x==xprobe][0] if x in xprobe else \
                               np.nan for x in xref ] for errs in yerrs ])

    return padded_yvals, padded_yerrs


def mkCurves1D(aos, ratio_panels, error_bars=True, variation_dict=None,
               colors=None, deviation=False, line_styles=['-', '--', '-.', ':'], **kwargs):

    """
    Auxiliary method that constructs and returns two text strings: one which represents
    the numerical data associated with the curves as well as edges, and one which
    contains the logic to plot these for a 1D plot.
    """

    cmd = yoda.util.StringCommand() # meta-data for top-level script
    data = yoda.util.StringCommand() # numerical values of the curves

    if isinstance(error_bars, bool):  # Convert to list of bool vals
        error_bars = [error_bars] * len(aos)

    if not colors:  # Use default mpl prop cycle vals
        colors = plt.rcParams['axes.prop_cycle'].by_key()['color']

    line_properties = LineProperties(colors, line_styles)

    for k,v in aos.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())
            aos[k] = scatter_2D

    # Print info about ref curve
    ref_curve = list(aos.values())[0]
    xpoints = 0.5*(ref_curve.xMins() + ref_curve.xMaxs())
    xedges = np.append(ref_curve.xMins(), ref_curve.xMax())
    data.add("from numpy import nan")
    data.newline()
    data.add(f"xpoints  = {[val for val in xpoints]}",
             f"xedges   = {[val for val in xedges]}",
             f"xmins    = {[val for val in ref_curve.xMins()]}",
             f"xmaxs    = {[val for val in ref_curve.xMaxs()]}")
    data.add(f"xerrs = [")
    data.add(f"  [abs(xpoints[i] - xmins[i])   for i in range(len(xpoints))],")
    data.add(f"  [abs(xmaxs[i]   - xpoints[i]) for i in range(len(xpoints))]")
    data.add(f"]")
    data.newline()

    cmd.newline()
    cmd.add("legend_handles = [] # keep track of handles for the legend")

    # Include values for curves
    yvals, yedges, yerrs, yups, ydowns, styles = {}, {}, {}, {}, {}, {}
    hasBand = {}
    variation_vals = {}
    band_vals, ratio_band_vals = {}, {}
    for i, (k, ao) in enumerate(aos.items()):

        colidx, linestyle = next(line_properties)
        zorder = 5 + i
        isRef = ao.annotation("IsRef", 0)

        ao_yvals = ao.yVals()
        ao_xvals = ao.xVals()
        ao_yerrs = np.array([ [err[0] for err in ao.yErrs()] if error_bars[i] else [0] * len(ao_yvals),
                             [err[1] for err in ao.yErrs()] if error_bars[i] else [0] * len(ao_yvals)
                   ])

        # in the case of a multi-energy histogram with missing values in the MC,
        # pad the missing values with nan
        if ao_xvals.shape != xpoints.shape:
            ao_yvals, ao_yerrs  = nanpad(ao_xvals, xpoints, ao_yvals, ao_yerrs)

        # save numerical values in dictionaries, write to file later
        # For "edges" we prepend the first element to be able to
        # draw a line like a rectangular histogram on canvas.
        yvals[ f"curve{i}"] = ao_yvals
        yedges[f"curve{i}"] = np.insert(ao_yvals, 0, ao_yvals[0])
        ydowns[f"curve{i}"] = ao_yerrs[0]
        yups[  f"curve{i}"] = ao_yerrs[1]
        styles[f"curve{i}"] = {
          'color' : f'colors[{colidx}]',
          'linestyle' : ao.annotation('LineStyle', linestyle),
          'marker' : ao.annotation('MarkerStyle', 'o'),
          'zorder' : zorder,
          'histstyle' : not ao.annotation('IsRef', 0),
        }

        # for every curve, save every variation curve to dict
        if variation_dict:

            var_set = variation_dict[k]
            for varKey, varAO in var_set.items():
                if varAO.type() == 'Scatter1D':
                    var_set[varKey] = yoda.core.Scatter2D()
                    var_set[varKey].addPoint(1.0, varAO.point(0).x(), 0.5, varAO.point(0).xErrs())
                var_yvals = varAO.yVals()
                if var_yvals.shape != xpoints.shape:
                    var_yvals, _ = nanpad(ao_xvals, xpoints, var_yvals)
                variation_vals[f"curve{i}_{varKey}"] = np.insert(var_yvals, 0, var_yvals[0])

            # this is a nominal Scatter2D with total PDF band as the
            # uncertainties, treat separately from the other multiweights
            if 'BandUncertainty' in var_set:
                hasBand[f'curve{i}'] = list()
                absoluteBandErrsDown, absoluteBandErrsUp = zip(*var_set['BandUncertainty'].yErrs())
                totalBandErrsDown = ao_yvals - absoluteBandErrsDown
                totalBandErrsUp = ao_yvals + absoluteBandErrsUp

                if not isRef:
                    totalBandErrsDown_ao = band_vals[f"curve{i}_band_dn"] =  np.insert(totalBandErrsDown, 0, totalBandErrsDown[0])
                    totalBandErrsUp_ao = band_vals[f"curve{i}_band_up"] = np.insert(totalBandErrsUp, 0, totalBandErrsUp[0])

                for r, (ratioLabel, ratioList) in enumerate(ratio_panels.items()):
                    hasBand[f'curve{i}'].append(f'ratio{r}')
                    ref_yvals = aos[ratioList[0]].yVals()
                    ratio_ref = np.insert(ref_yvals, 0, ref_yvals[0]) # pad for rectangular hist line

                    ratio_band_vals[f"ratio{r}_curve{i}_band_dn"] = np.divide(totalBandErrsDown_ao, ratio_ref,
                                                                     out=np.where(totalBandErrsDown_ao!=0, np.nan, 1),
                                                                     where=ratio_ref!=0)

                    ratio_band_vals[f"ratio{r}_curve{i}_band_up"] = np.divide(totalBandErrsUp_ao, ratio_ref,
                                                                     out=np.where(totalBandErrsUp_ao!=0, np.nan, 1),
                                                                     where=ratio_ref!=0)

                del var_set['BandUncertainty']


    # write dictionary of labels + yvalues to iterate over in executable python script
    writeLists(data, 'yvals',  yvals)
    writeLists(data, 'yedges', yedges)
    writeLists(data, 'yups',   yups)
    writeLists(data, 'ydowns', ydowns)

    # write out values for band and/or variation curves
    if hasBand:
        writeLists(data, 'band_edges', band_vals)
    if variation_dict:
        writeLists(data, 'variation_yvals', variation_vals)

    cmd.newline()
    cmd.add("# style options for curves",
            "# starts at zorder>=5 to draw curve on top of legend")

    # write out a dictionary that contains color, linestyle and zorder for every input YODA curve
    cmd.add("styles = {")
    for i, (key, val) in enumerate(styles.items()):
        cmd.add(f"  '{key}': " + "{")
        for ky, vl in val.items():
            if ky == 'color':
                cmd.add(f"""    '{ky}' : {vl.strip("'")},""")
            elif isinstance(vl, str):
                cmd.add(f"""    '{ky}' : '{vl}',""")
            else:
                cmd.add(f"""    '{ky}' : {vl},""")
        if i < len(styles) - 1:
            cmd.add("  },")
        else:
            cmd.add("  },", "}")
    cmd.add("# curve from input yoda files in main panel",
            "for label in dataf['yvals'].keys():",
            "    if all(np.isnan(v) for v in dataf['yvals'][label]):",
            "        continue",
            "    tmp = None",
            "    if styles[label]['histstyle']: # draw as histogram",
            "        tmp, = ax.plot(dataf['xedges'], dataf['yedges'][label],",
            "                       color=styles[label]['color'],",
            "                       linestyle=styles[label]['linestyle'],",
            "                       drawstyle='steps-pre', solid_joinstyle='miter',",
            "                       zorder=styles[label]['zorder'], label=label)",
            "        ax.errorbar(dataf['xpoints'], dataf['yvals'][label], color=styles[label]['color'],",
            "                    xerr=dataf['xerrs'], yerr=[dataf['ydowns'][label], dataf['yups'][label]],",
            "                    linestyle='none', zorder=styles[label]['zorder'])",
           "    else: # draw as scatter",
            "        tmp = ax.errorbar(dataf['xpoints'], dataf['yvals'][label],",
            "                          xerr=dataf['xerrs'],",
            "                          yerr=[ dataf['ydowns'][label],",
            "                          dataf['yups'][label] ],",
            "                          fmt=styles[label]['marker'],",
            "                          ecolor=styles[label]['color'],",
            "                          color=styles[label]['color'], zorder=2)",
            "        tmp[-1][0].set_linestyle(styles[label]['linestyle'])",
            "    legend_handles += [tmp]")

    if hasBand:
        opac = ao.annotation('ErrorBandOpacity', 0.2)
        # make sure every variation matches color of nominal value
        cmd.add(f"    if dataf['band_edges'].get(label+'_band_dn', None):",
                f"        ax.fill_between(dataf['xedges'],",
                f"                        dataf['band_edges'][label+'_band_dn'], dataf['band_edges'][label+'_band_up'],",
                f"                        color=styles[label]['color'], alpha={opac},step='pre',",
                f"                        zorder=styles[label]['zorder'],edgecolor=None)")
    elif variation_dict:
        # make sure every variation (multiweight) matches color of nominal value
        cmd.add(f"    for varLabel in dataf['variation_yvals'].keys():",
                f"        if varLabel.startswith(label):",
                f"            tmp, = ax.plot(dataf['xedges'], dataf['variation_yvals'][varLabel],",
                f"                           color=styles[label]['color'],",
                f"                           linestyle=styles[label]['linestyle'],",
                f"                           drawstyle='steps-pre', solid_joinstyle='miter',",
                f"                           zorder=styles[label]['zorder'], alpha=0.5)")


    # Now loop over ratio panels and add curves per panel
    for i, (ratioLabel, ratioList) in enumerate(ratio_panels.items()):

        prefix  = 'ratio%d' % i
        axlabel = prefix+'_ax'

        data.add("\n\n# lists for ratio plot")
        cmd.add("\n\n# plots on ratio panel")

        if isinstance(error_bars, bool):  # Convert to list of bool vals
            error_bars = [error_bars] * len(ratioList)

        if not colors:  # Use default mpl prop cycle vals
            colors = plt.rcParams['axes.prop_cycle'].by_key()['color']

        line_properties = LineProperties(colors, line_styles)

        # Reference of current ratio panel
        ref_yvals = aos[ratioList[0]].yVals()
        ref_yerrs = aos[ratioList[0]].yErrs()

        # Now add non-reference curves to current ratio panel
        ratio_yvals, ratio_yups, ratio_ydowns = {}, {}, {}
        ratio_variation_vals, ratio_bands = {}, {}

        # Plot the ratio curves
        for idx, k in enumerate(ratioList):

            ao = aos[k]
            isRef = ao.annotation("IsRef", 0)

            # values for central curve
            ao_xvals = ao.xVals()
            ao_yvals = ao.yVals()
            ao_yerrs = np.array([ [err[0] for err in ao.yErrs()] if error_bars[i] else [0] * len(ao_yvals),
                                  [err[1] for err in ao.yErrs()] if error_bars[i] else [0] * len(ao_yvals)
                                ])
            colidx, linestyle = next(line_properties)

            # in the case of a multi-energy histogram with missing values in the MC,
            # pad the missing values with nan
            if ao_xvals.shape != xpoints.shape:
                ao_yvals, ao_yerrs = nanpad(ao_xvals, xpoints, ao_yvals, ao_yerrs)

            num_yvals = ao_yvals
            den_yvals = ref_yvals
            if not isRef:
                num_yvals = np.insert(num_yvals, 0, num_yvals[0])
                den_yvals = np.insert(den_yvals, 0, den_yvals[0])

            # set ratio to 1 if we are dividing 0 by 0, nan if dividing x by 0
            ratio = np.divide(num_yvals, den_yvals,
                              out=np.where(num_yvals!=0, np.nan, 1),
                              where=den_yvals!=0)

            # values for uncertainties on central curve
            if error_bars[idx]:

                num_errminus = ao_yvals - ao_yerrs[0]
                num_errplus  = ao_yvals + ao_yerrs[1]

                # construct ratios for up/down errors
                # set ratio to 1 if we are dividing 0 by 0, nan if dividing x by 0
                ratio_ydowns[f"curve{idx}"] = np.divide(num_errminus, ref_yvals,
                                                      out=np.where(num_errminus!=0, np.nan, 1),
                                                      where=ref_yvals!=0)
                ratio_yups[f"curve{idx}"] = np.divide(num_errplus, ref_yvals,
                                                    out=np.where(num_errplus!=0, np.nan, 1),
                                                    where=ref_yvals!=0)
            else:
                ratio_ydowns[f"curve{idx}"] = ratio_yups[f'curve{idx}'] = [0] * ref_yvals # results in no error bars drawn



            if k in ratio_bands:
                totalBandErrsDown_ao = band_vals[f"curve{i}_band_dn"]
                totalBandErrsUp_ao = band_vals[f"curve{i}_band_up"]

                band_vals[f"{prefix}_curve{idx}_band_dn"] = np.divide(totalBandErrsDown_ao, ref_yvals,
                                                                   out=np.where(totalBandErrsDown_ao!=0, np.nan, 1),
                                                                   where=ref_yvals!=0)

                band_vals[f"{prefix}_curve{idx}_band_up"] = np.divide(totalBandErrsUp_ao, ref_yvals,
                                                                   out=np.where(totalBandErrsUp_ao!=0, np.nan, 1),
                                                                   where=ref_yvals!=0)
            elif k in variation_dict and not k in ratio_bands:
                # values for variation curves, skip when combined band is already plotted
                var_set = variation_dict[k]
                for varKey, varAO in var_set.items():
                    var_yvals = varAO.yVals()
                    if var_yvals.shape != ref_yvals.shape:
                        var_yvals, _ = nanpad(ao_xvals, xpoints, var_yvals)
                    vals_mw = np.insert(var_yvals, 0, var_yvals[0])

                    ratio_mw = np.divide(vals_mw, den_yvals,
                                         out=np.where(vals_mw!=0, np.nan, 1),
                                         where=den_yvals!=0)
                    ratio_variation_vals[f"curve{idx}_{varKey}"] = ratio_mw

            if deviation:
                # average up and down uncertainties
                # TODO decide on average or max value for (possibly) asymmetric error bars
                ref_yerrs_avg = np.array([ 0.5 * (err[0] + err[1]) for err in ref_yerrs])

                # use the band uncertainty for the curve if specified by user, other just stat. uncertainty
                if k in ratio_bands:
                    totalBandErrsDown_ao = band_vals[f"curve{i}_band_dn"]
                    totalBandErrsUp_ao = band_vals[f"curve{i}_band_up"]
                    ao_yerrs_avg = 0.5 * np.array(totalBandErrsDown_ao + totalBandErrsUp_ao)
                else:
                    ao_yerrs_avg = 0.5 * (ao_yerrs[0,:] + ao_yerrs[1,:])

                # nominal values have length nbins + 1, whereas errors have length nbins
                ratio = num_yvals - den_yvals
                if not isRef:
                    ratio = ratio[1:]
                ratio /= np.sqrt(ref_yerrs_avg ** 2 + ao_yerrs_avg **2 )
                ratio = np.insert(ratio, 0, ratio[0])

            ratio_yvals[f"curve{idx}"] = ratio

        # write lists to read in executable py script
        writeLists(data, f'{prefix}_yvals',  ratio_yvals)
        writeLists(data, f'{prefix}_yups',   ratio_yups)
        writeLists(data, f'{prefix}_ydowns', ratio_ydowns)

        data.add(f"{prefix}_yerrs = {{")
        for label, vals in ratio_yvals.items():
            offset = len(vals) - len(ratio_ydowns[label])
            lo = np.abs(np.subtract(vals[offset:], ratio_ydowns[label]))
            hi = np.abs(np.subtract(ratio_yups[label], vals[offset:]))
            data.add(f"""  '{label}' : [""",
                     f"""    {[val for val in lo]},""",
                     f"""    {[val for val in hi]},""",
                     f"""  ],""")
        data.add("}")

        if variation_dict and not k in ratio_bands:
            writeLists(data, f'{prefix}_variation_vals', ratio_variation_vals)

        # for loop to plot curve in ratio
        cmd.add(f"# curve from input yoda files in ratio panel",
                f"for label, yvals in dataf['{prefix}_yvals'].items():",
                f"    if all(np.isnan(v) for v in yvals):",
                f"        continue")
        indent = ''
        if not deviation:
            indent += '    '
            cmd.add(f"    if styles[label]['histstyle']: # plot as histogram")
        cmd.add(f"{indent}    {axlabel}.plot(dataf['xedges'], yvals,",
                f"{indent}                  color=styles[label]['color'],",
                f"{indent}                  linestyle=styles[label]['linestyle'],",
                f"{indent}                  drawstyle='steps-pre', zorder=1,",
                f"{indent}                  solid_joinstyle='miter')")
        if not deviation:
            cmd.add(f"        {axlabel}.vlines(dataf['xpoints'],",
                    f"                         dataf['{prefix}_ydowns'][label], dataf['{prefix}_yups'][label],",
                    f"                         color=styles[label]['color'], zorder=1)",
                    f"    else: # plot as scatter",
                    f"        tmp = {axlabel}.errorbar(dataf['xpoints'],",
                    f"                                 yvals, xerr=dataf['xerrs'],",
                    f"                                 yerr=dataf['{prefix}_yerrs'][label],",
                    f"                                 fmt=styles[label]['marker'],",
                    f"                                 ecolor=styles[label]['color'],",
                    f"                                 color=styles[label]['color'])",
                    f"        tmp[-1][0].set_linestyle(styles[label]['linestyle'])")
        cmd.newline()

        if deviation:
            # fill area in deviation mode
            cmd.add(f"# coloured area displaying 1-sigma region and horizontal in ratio panel",
                    f"{axlabel}.fill_between(dataf['xedges'], [-1.] * len(dataf['xedges']),",
                    f"  [1.] * len(dataf['xedges']), color='yellow', alpha=0.2, zorder=1)",
                    f"{axlabel}.plot(dataf['xedges'], [0.] * len(dataf['xedges']),",
                    f"                  color='black', zorder=2, linewidth=0.5)")
        else:
            # for loop to plot variations in ratio
            if hasBand:
                opac = ao.annotation('ErrorBandOpacity', 0.2)
                cmd.add(f"    if dataf['ratio_band_edges'].get('{prefix}_'+label+'_band_dn', None):",
                        f"        {axlabel}.fill_between(dataf['xedges'],",
                        f"                              dataf['ratio_band_edges']['{prefix}_'+label+'_band_dn'],",
                        f"                              dataf['ratio_band_edges']['{prefix}_'+label+'_band_up'],",
                        f"                              color=styles[label]['color'],",
                        f"                              alpha={opac},step='pre',",
                        f"                              zorder=styles[label]['zorder'],edgecolor=None)")

            if k in variation_dict:
                cmd.add(f"    for varlabel in dataf['{prefix}_variation_vals'].keys():",
                        f"        if varlabel.startswith(label):",
                        f"            {axlabel}.plot(dataf['xedges'],",
                        f"                            dataf['{prefix}_variation_vals'][varlabel],",
                        f"                            color=styles[label]['color'],",
                        f"                            linestyle=styles[label]['linestyle'],",
                        f"                            drawstyle='steps-pre', solid_joinstyle='miter',",
                        f"                            zorder=1, alpha=0.5)")



        cmd.newline()
    if hasBand:
        writeLists(data, 'ratio_band_edges', ratio_band_vals)

    return cmd.get(), data.get()


def writeLists(cmd, val_type, val_dict):
    """Take a label and dictionary as input,
       return as a formatted string."""
    cmd.add(f"""{val_type} = {{""")
    for label, vals in val_dict.items():
        cmd.add(f"""  '{label}' : {[val for val in vals]},""")
    cmd.add("}")


class LineProperties:

    def __init__(self, colors, linestyles):
        self.colors = colors
        self.linestyles = linestyles
        self.color_index = 0
        self.style_index = 0

    def __iter__(self):
        return self

    def __next__(self):
        vals = (self.color_index, self.linestyles[self.style_index])
        self.color_index += 1
        if self.color_index == len(self.colors):
            self.color_index = 0
            self.style_index += 1
            if self.style_index == len(self.linestyles):
                self.style_index = 0
        return vals
