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

import yoda
import yoda.plotting.utils as putils

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(), max(aos[0].xMaxs()))
    ref_xvals = aos[0].xVals()
    ref_xerrs = aos[0].xErrs()
    ypoints   = 0.5*(aos[0].yMins() + aos[0].yMaxs())
    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(), max(aos[0].zMaxs()))
    ref_zvals = aos[0].zVals()
    ref_zerrs = aos[0].zErrs()
    data.add(f"import numpy as np",
             f"from numpy import nan",
             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"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 = dict() # 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['ref'] = 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()]
            c_zdowns[f'c{i}'] = [err[0] for err in ao.zErrs()]
            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[label] = tmp_contour")

    cmd.newline()

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


def mkCurves1D(aos, ratio_panels, error_bars=True, variation_dict=None,
               colors=None, deviation_modes=True, 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 isinstance(deviation_modes, bool):  # Convert to list of bool vals
        deviation_modes = [deviation_modes] * len(ratio_panels)

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

    line_properties = LineProperties(colors, line_styles)

    ref_label = list(aos.keys())[0]
    ref_curve = list(aos.values())[0]
    ref_xpoints = 0.5*(ref_curve.xMins() + ref_curve.xMaxs()) # forces bin centre
    ref_xedges = np.append(ref_curve.xMins(), max(ref_curve.xMaxs()))

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

    # Print values for dependent axes
    xpoints, xedges, xups, xdowns = {}, {}, {}, {}
    yvals, yups, ydowns = {}, {}, {}
    band_vals, ratio_band_vals, hasBand = {}, {}, { 'main' : [] }
    variation_vals, styles = {}, {}
    for i, (k, ao) in enumerate(aos.items()):

        colidx, linestyle = next(line_properties)
        zorder = 5 + i

        xpoints[k] = 0.5*(ao.xMins() + ao.xMaxs()) # forces bin centre
        xedges[ k] = np.append(ao.xMins(), max(ao.xMaxs()))

        # in the case of a multi-energy histogram with missing values in the MC, pad missing values with nan
        ao_yvals, ao_xerrs, ao_yerrs = putils.reshape(ref_xpoints, ref_xedges, ao)

        # 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[ k] = ao_yvals
        xdowns[k], xups[k] = np.abs(ao_xerrs)
        ydowns[k], yups[k] = np.abs(ao_yerrs)

        # check if the AO needs an uncertainty band
        if str(ao.annotation('ErrorBand', '0')) != '0':
            # This AO wants the error bars displayed as a band
            if variation_dict is None:  variation_dict = { k : {} }
            if 'BandUncertainty' not in variation_dict[k]:
                variation_dict[k]['BandUncertainty'] = ao.clone()

        # Save variations for current curve
        if variation_dict:

            var_set = variation_dict[k]
            for varKey, varAO in var_set.items():
                if varKey == 'BandUncertainty':  continue
                var_yvals, _, _ = putils.reshape(ref_xpoints, ref_xedges, varAO)
                variation_vals[k + f"_{varKey}"] = np.insert(var_yvals, 0, var_yvals[0])

            if any(deviation_modes) or 'BandUncertainty' in var_set:
                levels = [ f'{sig}sigma' \
                           for sig in reversed(str(ao.annotation('ErrorBand', '1')).split()) \
                           if int(sig) ]
                if 'BandUncertainty' in var_set:
                    # This curve has an uncertainty band in the main panel!
                    if len(levels) > len(hasBand['main']):
                        hasBand['main'] = levels

                    # account for multi-energy histos
                    var_yvals, var_xerrs, var_yerrs = putils.reshape(ref_xpoints, ref_xedges, var_set['BandUncertainty'])
                    for sigma in levels:
                        var_dn = var_yvals - np.abs(var_yerrs[0]) * float(sigma[:-5])
                        var_up = var_yvals + np.abs(var_yerrs[1]) * float(sigma[:-5])
                        var_dn = np.insert(var_dn, 0, var_dn[0])
                        var_up = np.insert(var_up, 0, var_up[0])
                        band_vals[k+f"_{sigma}-band_dn"] = var_dn
                        band_vals[k+f"_{sigma}-band_up"] = var_up

                # account for case that MC curve has finer binning than ref data
                bandao = var_set['BandUncertainty'] if 'BandUncertainty' in var_set else ao
                var_yvals, var_xerrs, var_yerrs = putils.reshape(ref_xpoints, ref_xedges, bandao, True)

                for r, (ratioLabel, ratioList) in enumerate(ratio_panels.items()):
                    if 'BandUncertainty' not in var_set and not deviation_modes[r]:  continue
                    rshort = f'ratio{r}'
                    den_yvals, _, _ = putils.reshape(ref_xpoints, ref_xedges, aos[ratioList[0]], True)
                    den_yvals = np.insert(den_yvals, 0, den_yvals[0])
                    levels = [ f'{sig}sigma' \
                               for sig in reversed(str(ao.annotation(ratioLabel+'ErrorBand',
                                                       ao.annotation('ErrorBand', '1'))).split()) \
                               if int(sig) ]
                    if len(levels) > len(hasBand.get(rshort, [])):
                        hasBand[rshort] = levels
                    for sigma in levels:
                        pre = f"{rshort}_{k}_{sigma}-band"
                        if deviation_modes[r]:
                            ref_yerrs_avg = ref_curve.errAvgs(1)
                            var_yerrs_avg = 0.5*np.abs(var_yerrs[0] + var_yerrs[1]) * float(sigma[:-5])
                            offset = var_yvals - den_yvals[1:]
                            var_dn = putils.safeDiv(offset - var_yerrs_avg, ref_yerrs_avg)
                            var_up = putils.safeDiv(offset + var_yerrs_avg, ref_yerrs_avg)
                            ratio_band_vals[pre+f"_dn"] = np.insert(var_dn, 0, var_dn[0])
                            ratio_band_vals[pre+f"_up"] = np.insert(var_up, 0, var_up[0])
                        else:
                            var_dn = var_yvals - np.abs(var_yerrs[0]) * float(sigma[:-5])
                            var_up = var_yvals + np.abs(var_yerrs[1]) * float(sigma[:-5])
                            ratio_band_vals[pre+f"_dn"] = putils.safeDiv(np.insert(var_dn, 0, var_dn[0]), den_yvals)
                            ratio_band_vals[pre+f"_up"] = putils.safeDiv(np.insert(var_up, 0, var_up[0]), den_yvals)

                if 'BandUncertainty' in var_set:
                    del var_set['BandUncertainty']

        # set cosmetic defaults
        styles[k] = { 'color' : colors[colidx] }
        for r, pre in enumerate(['']+list(ratio_panels.keys())):
            rlabel = 'ratio%d_' % (r-1) if r else ''
            hasBand.setdefault(rlabel[:-1], [])
            histstyle = ao.annotation(pre+'ConnectBins', styles[k]['histstyle'] if pre else 1)
            lstyle = ao.annotation(pre+'LineStyle', ao.annotation('LineStyle', linestyle))
            lwidth = ao.annotation(pre+'LineWidth', ao.annotation('LineWidth', 1))
            mstyle = ao.annotation(pre+'MarkerStyle', styles[k]['marker'] if pre else 'none' if histstyle else 'o')
            msize = ao.annotation(pre+'MarkerSize', ao.annotation('MarkerSize', 2))
            capsize = ao.annotation(pre+'ErrorCapSize', ao.annotation('ErrorCapSize', 0.0))
            xerrbars = ao.annotation(pre+'ErrorBars', ao.annotation('ErrorBars', int(error_bars[i])))
            default_yerrbar = int(error_bars[i] and (not deviation_modes[r-1] if pre else True))
            yerrbars = ao.annotation(pre+'ErrorBars', ao.annotation('ErrorBars', default_yerrbar))
            dstyle = None if ao.annotation(pre+'ConnectMarkers', ao.annotation('ConnectMarkers', 0)) else 'steps-pre'
            styles[k].update({
                rlabel+'linestyle'   : lstyle,
                rlabel+'linewidth'   : lwidth,
                rlabel+'marker'      : mstyle,
                rlabel+'markersize'  : msize,
                rlabel+'capsize'     : capsize,
                rlabel+'zorder'      : zorder,
                rlabel+'histstyle'   : histstyle,
                rlabel+'drawstyle'   : dstyle,
                rlabel+'xerrorbars'  : int(xerrbars),
                rlabel+'yerrorbars'  : int(yerrbars),
            })

            bandcol = ao.annotation(pre+'ErrorBandColor', ao.annotation('ErrorBandColor', colors[colidx])).split()
            bandstyle = ao.annotation(pre+'ErrorBandStyle', ao.annotation('ErrorBandStyle', None))
            bandstyle = [None] if bandstyle is None else bandstyle.split()
            bandopac = list(ao.annotation(pre+'ErrorBandOpacity', ao.annotation('ErrorBandOpacity', [0.2])))
            for s, sigma in enumerate(hasBand[rlabel[:-1] if pre else 'main']):
                styles[k].update({
                    rlabel+sigma+'bandcolor'   : bandcol[0] if len(bandcol) == 1 else bandcol[-1-s],
                    rlabel+sigma+'bandstyle'   : bandstyle[0] if len(bandstyle)==1 else bandstyle[-1-s],
                    rlabel+sigma+'bandopacity' : bandopac[0] if len(bandopac)==1 else bandopacity[-1-s],
                })
            if pre == '':
                styles[k].update({
                    'fillcolor'   : ao.annotation('FillColor', None),
                    'fillopacity' : ao.annotation('FillOpacity', 1.0),
                })



    # write dictionary of labels + yvalues to iterate over in executable python script
    data.add("import numpy as np",
             "from numpy import nan")
    data.newline()
    data.add("add_legend_handle = [",
             ",\n".join([ f"    '{k}'" for k,v in aos.items() if v.annotation('Title', None) ]),
             "]")
    data.newline()
    writeLists(data, 'xpoints', xpoints)
    writeLists(data, 'xedges',  xedges)
    data.add(f"ref_xerrs = [")
    data.add(f"  [abs(xpoints['{ref_label}'][i]   - xedges['{ref_label}'][i])"
             + f" for i in range(len(xpoints['{ref_label}']))],")
    data.add(f"  [abs(xedges['{ref_label}'][i+1] - xpoints['{ref_label}'][i])"
             + f" for i in range(len(xpoints['{ref_label}']))]")
    data.add(f"]")
    data.newline()
    writeLists(data, 'yvals',  yvals)
    writeLists(data, 'xerrs', xups, xdowns)
    writeLists(data, 'yerrs', yups, ydowns)

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

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

    # add logic to draw curves in main panel
    cmd.add(f"# curve from input yoda files in main panel",
            f"for label, yvals in dataf['yvals'].items():",
            f"    if all(np.isnan(v) for v in dataf['yvals'][label]):",
            f"        continue",
            f"    tmp = None",
            f"    if styles[label]['histstyle']: # draw as histogram",
            f"        xpos = dataf['xedges'][label] if styles[label]['drawstyle'] else dataf['xpoints'][label]",
            f"        ypos = np.insert(yvals, 0, yvals[0]) if styles[label]['drawstyle'] else yvals",
            f"        if styles[label]['fillcolor']: # fill area below curve",
            f"            ax.fill_between(xpos, ypos, step='pre',",
            f"                            color=styles[label]['fillcolor'],",
            f"                            alpha=styles[label]['fillopacity'])",
            f"        tmp, = ax.plot(xpos, ypos,",
            f"                       color=styles[label]['color'],",
            f"                       linestyle=styles[label]['linestyle'],",
            f"                       linewidth=styles[label]['linewidth'],",
            f"                       drawstyle=styles[label]['drawstyle'], solid_joinstyle='miter',",
            f"                       zorder=styles[label]['zorder'], label=label)",
            f"    tmp = ax.errorbar(dataf['xpoints'][label], yvals,",
            f"                      xerr=np.array(dataf['xerrs'][label])*styles[label]['xerrorbars'],",
            f"                      yerr=np.array(dataf['yerrs'][label])*styles[label]['yerrorbars'],",
            f"                      fmt=styles[label]['marker'], capsize=styles[label]['capsize'],",
            f"                      markersize=styles[label]['markersize'],",
            f"                      ecolor=styles[label]['color'],",
            f"                      color=styles[label]['color'], zorder=styles[label]['zorder'])",
            f"    tmp[-1][0].set_linestyle(styles[label]['linestyle'])",
            f"    tmp[-1][0].set_linewidth(styles[label]['linewidth'])",
            f"    if label in dataf['add_legend_handle']:  legend_handles[label] = tmp")

    for sigma in hasBand['main']:
        cmd.add(f"    if dataf['band_edges'].get(label+'_{sigma}-band_dn', None):",
                f"        ax.fill_between(dataf['xedges'][label],",
                f"                        dataf['band_edges'][label+'_{sigma}-band_dn'],",
                f"                        dataf['band_edges'][label+'_{sigma}-band_up'],",
                f"                        color=styles[label]['{sigma}bandcolor'],",
                f"                        alpha=styles[label]['{sigma}bandopacity'],",
                f"                        hatch=styles[label]['{sigma}bandstyle'], step='pre',",
                f"                        zorder=styles[label]['zorder'], edgecolor=None)")
    if 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'][label], dataf['variation_yvals'][varLabel],",
                f"                           color=styles[label]['color'],",
                f"                           linestyle=styles[label]['linestyle'],",
                f"                           linewidth=styles[label]['linewidth'],",
                f"                           drawstyle='steps-pre', solid_joinstyle='miter',",
                f"                           zorder=styles[label]['zorder'], alpha=0.5)")


    # Now loop over ratio panels and add respective curves
    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 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, _, _ = putils.reshape(ref_xpoints, ref_xedges, aos[ratioList[0]], True)

        # Plot the curves in the current ratio panel
        ratio_yvals, ratio_yups, ratio_ydowns = {}, {}, {}
        ratio_variation_vals = {}

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

            ao = aos[k]
            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
            ao_yvals, ao_xerrs, ao_yerrs = putils.reshape(ref_xpoints, ref_xedges, ao, True)

            # set ratio to 1 if we are dividing 0 by 0, nan if dividing x by 0
            ratio = putils.safeDiv(ao_yvals, ref_yvals)
            ratio_ydowns[k] = np.abs(putils.safeDiv(ao_yerrs[0], ref_yvals))
            ratio_yups[k] = np.abs(putils.safeDiv(ao_yerrs[1], ref_yvals))

            if k in variation_dict:
                # 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, _, _ = putils.reshape(ref_xpoints, ref_xedges, varAO, True)
                    ratio_mw = putils.safeDiv(var_yvals, ref_yvals)
                    ratio_variation_vals[k + f"_{varKey}"] = np.insert(ratio_mw, 0, ratio_mw[0])

            if deviation_modes[i]:
                # average up and down uncertainties of reference
                ref_yerrs_avg = ref_curve.errAvgs(1)
                ratio = ao_yvals - ref_yvals
                ratio = putils.safeDiv(ratio, ref_yerrs_avg)
                ao_yerrs_avg = 0.5*(ao_yerrs[0]+ao_yerrs[1])
                ratio_yups[k] = ratio_ydowns[k] = np.abs(putils.safeDiv(ao_yerrs_avg, ref_yerrs_avg))

            ratio_yvals[k] = ratio

        # write lists to read in executable py script
        writeLists(data, f'{prefix}_yvals',  ratio_yvals)
        writeLists(data, f'{prefix}_yerrs', ratio_yups, ratio_ydowns)
        if variation_dict:
            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",
                f"    if styles[label]['{prefix}_histstyle']: # plot as histogram",
                f"        xpos = dataf['xedges']['{ref_label}'] " \
                + f"if styles[label]['{prefix}_drawstyle'] else dataf['xpoints']['{ref_label}']",
                f"        ypos = np.insert(yvals, 0, yvals[0]) if styles[label]['{prefix}_drawstyle'] else yvals",
                f"        {axlabel}.plot(xpos, ypos,",
                f"                       color=styles[label]['color'],",
                f"                       linewidth=styles[label]['linewidth'],",
                f"                       linestyle=styles[label]['{prefix}_linestyle'],",
                f"                       drawstyle=styles[label]['{prefix}_drawstyle'], zorder=styles[label]['zorder'],",
                f"                       solid_joinstyle='miter')",
                f"    tmp = {axlabel}.errorbar(dataf['xpoints']['{ref_label}'], yvals,",
                f"                             xerr=np.array(dataf['ref_xerrs'])*styles[label]['{prefix}_xerrorbars'],",
                f"                             yerr=np.array(dataf['{prefix}_yerrs'][label])*styles[label]['{prefix}_yerrorbars'],",
                f"                             fmt=styles[label]['{prefix}_marker'], capsize=styles[label]['{prefix}_capsize'],",
                f"                             markersize=styles[label]['{prefix}_markersize'],",
                f"                             ecolor=styles[label]['color'],",
                f"                             color=styles[label]['color'])",
                f"    tmp[-1][0].set_linestyle(styles[label]['{prefix}_linestyle'])",
                f"    tmp[-1][0].set_linewidth(styles[label]['{prefix}_linewidth'])")
        cmd.newline()

        # for loop to plot variations in ratio
        for sigma in hasBand[prefix]:
            cmd.add(f"    if dataf['ratio_band_edges'].get('{prefix}_'+label+'_{sigma}-band_dn', None):",
                    f"        {axlabel}.fill_between(dataf['xedges']['{ref_label}'],",
                    f"                              dataf['ratio_band_edges']['{prefix}_'+label+'_{sigma}-band_dn'],",
                    f"                              dataf['ratio_band_edges']['{prefix}_'+label+'_{sigma}-band_up'],",
                    f"                              color=styles[label]['{prefix}_{sigma}bandcolor'],",
                    f"                              hatch=styles[label]['{prefix}_{sigma}bandstyle'],",
                    f"                              alpha=styles[label]['{prefix}_{sigma}bandopacity'],",
                    f"                              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']['{ref_label}'],",
                    f"                            dataf['{prefix}_variation_vals'][varlabel],",
                    f"                            color=styles[label]['color'],",
                    f"                            linestyle=styles[label]['linestyle'],",
                    f"                            linewidth=styles[label]['linewidth'],",
                    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, val_alt = None):
    """Take a label and dictionary as input,
       return as a formatted string."""
    cmd.add(f"""{val_type} = {{""")
    for label, vals in val_dict.items():
        if val_alt is None:
            cmd.add(f"""    '{label}' : {[float(val) for val in vals]},""")
        else:
            cmd.add(f"""    '{label}' : [""",
                    f"""        {[float(val) for val in vals]},""",
                    f"""        {[float(val) for val in val_alt[label]]},""",
                    f"""    ],""")
    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
