#! /usr/bin/python3

"""\
%(prog)s [-d <refdir>|-R] [-w <weightfile>] [<ipolfiles>=ipol.dat ...] [-r <runsdir>=<refdir>/../mc] [args]

Use the interpolations stored in <ipolfile> to find optimised parameters with
the reference histograms found in the <refdir> as the optimisation target.

The <runsdir> is used to calculate the maximum error value seen for each bin,
to regularise interpolated errors which could otherwise blow up, leading to
an unrepresentative small chi2 (and hence fit result) outside the sampled ranges.


WEIGHT FILE SYNTAX:

The weight file syntax is derived from YODA path syntax, and allows selecting bin
ranges either by physical value x (identifying the bin it lies in) or by bin index n:

Whole histo:    /path/parts/to/histo            weight  {norm1/normA/normS}
Single bin:     /path/parts/to/histo@x          weight
Bin range:      /path/parts/to/histo@xmin:xmax  weight
Range from min: /path/parts/to/histo@xmin:      weight
Range to max:   /path/parts/to/histo@:xmax      weight
Single bin:     /path/parts/to/histo#n          weight
Bin range:      /path/parts/to/histo#nmin:nmax  weight
Range from min: /path/parts/to/histo#nmin:      weight
Range to max:   /path/parts/to/histo#:nmax      weight

Blank lines and lines starting with a # symbol will be ignored.

The bin indices used with the # syntax start at 0, and the end index
in a range is non-inclusive. In the range form, if xmin/nmin or
xmax/nmax is left blank, it accepts all bins from the start of the
histogram, or all bins to the end of the histogram respectively.

The norm1,normA,normS options on the whole-histogram selection
determine whether the histogram should be normalized to the reference
data in some way, resulting in a shape-only fit.  Respectively, these
fix the prediction's normalization to match the data on the first bin,
the area integral, or the sum of bin values.


TODO:
 * Use explicit param-name mapping rather than inspection when calling the GoF function
 * Restrict scanning to a hypersphere cf. prof2-sample
 * Include correlations in the tuning and resampling.
 * Handle run combination file/string (write a hash of the run list into the ipol filename?)

"""

import argparse, os, sys
ap = argparse.ArgumentParser(usage=__doc__,formatter_class=argparse.ArgumentDefaultsHelpFormatter)
ap.add_argument("IFILES", nargs="*", default=["ipol.dat"], help="The ipol file(s)")
ap.add_argument("-d", "--datadir", "--refdir", dest="REFDIR", default=None, help="The ref data directory")
ap.add_argument("-w", "--wfile", dest="WFILE", default=None, help="Path to a weight file to specify chi2 weights of each bin in the fit (default: %(default)s)")
ap.add_argument("-r", "--runsdir", dest="RUNSDIR", default=None, help="The directory containing MC runs (can be used to inform param bounds)")
ap.add_argument("-R", "--rivet", dest="USE_RIVET_REF", action="store_true", default=False, help="Use Rivet's API to fetch its reference path")
ap.add_argument("-o", "--outdir", dest="OUTDIR", default="tunes", help="Prefix for outputs ")
ap.add_argument("--limits", dest="LIMITS", default=None, help="Simple text file with parameter limits and fixed parameters")
ap.add_argument("-m", "--minos", dest="MINOS", default=False, action="store_true", help="Run Minos algorithm after minimisation")
ap.add_argument("--limit-errs", dest="USE_RUNSDIR", action="store_true", default=False, help="Re-read the runsdir to regularise error ipols")
ap.add_argument("--norm-weights", dest="NORMW", action="store_true", default=False, help="Divide observable weight by number of regular bins")
ap.add_argument("-x", "--allow-extrapolation", dest="EXTRAPOLATE", action="store_true", default=False, help="Allow the minimiser to go into region of extrapolation")
ap.add_argument("-v", "--debug", dest="DEBUG", action="store_true", default=False, help="Turn on some debug messages")
ap.add_argument("-q", "--quiet", dest="QUIET", action="store_true", default=False, help="Turn off messages")
ap.add_argument("-s", "--strategy", dest="STRATEGY",  default=2, type=int, help="Set Minuit strategy [0 fast, 1 standard, 2 slow]")
ap.add_argument("--filter", dest="FILTER", action="store_true", default=False, help="Filter out data bins that have 0 error")
ap.add_argument("--profiles", dest="PROFILES", action="store_true", default=False, help="Draw Minos and MIGRAD profiles")
ap.add_argument("--contours", dest="CONTOURS", action="store_true", default=False, help="Draw Minos and MIGRAD contours")
ap.add_argument("--max-attempts", dest="NUM_ATTEMPTS", default=10, type=int, help="Max number of attempts to reach convergence")
ap.add_argument("--max-calls", dest="NUM_CALLS", default=10000, type=int, help="Max number of iterations allowed in an attempt")
ap.add_argument("--scan-num", dest="SCANNP", default=0, type=int, help="Number of test points to find a good MIGRAD start point")
ap.add_argument("--scan-sampler", dest="SCANSAMPLER", default="uniform", help="Sampler to use for start point finding uniform|latin|sobol")
ap.add_argument("--scan-seed", dest="SCANSEED", default=0, type=int , help="Scan-sampler seed ")
ap.add_argument("--gof-function", dest="GOFFUNCTION", default="chi2", help="Specify the goodness-of-fit function. Options are chi2, chiRegularized, tanhChi2, chi2xsec, sigmoidChi2")
ap.add_argument("--theory-error", dest="THEORYERROR", default=0.0, type=float, help="Specify a theory error in percent")
args = ap.parse_args()


## Sort out conflicts in ref-data source
if args.USE_RIVET_REF and args.REFDIR:
    print("Error: you can only specify one of the -d/--datadir and -R/--rivet flags")
    sys.exit(1)

## Try to get the ref-data dir from Rivet
if args.USE_RIVET_REF:
    try:
        import rivet
        rivetpaths = rivet.getAnalysisDataPaths()
        if len(rivetpaths) > 1:
            print(f"Warning: Rivet provided multiple ref data paths, using only the first: {rivetpaths[0]}")
        args.REFDIR = rivetpaths[0]
    except ImportError:
        print(f"Warning: Rivet ref data requested, but module could not be imported")
        args.USE_RIVET_REF = False
    except IndexError:
        print(f"Warning: Rivet ref data requested, but module provided an empty path")
        args.USE_RIVET_REF = False

## Exit with an error if there is still no ref-data source
if not args.REFDIR:
    print("Error: no ref-data path found, either directly or from Rivet")
    sys.exit(1)


## Check availability of a minimiser
try:
    from iminuit import Minuit
    import numpy as np
except ImportError as e:
    print("Unable to import iminuit, exiting", e)
    print("Try installing iminuit with pip: pip install iminuit")
    sys.exit(1)


# TODO: clean up
assert args.GOFFUNCTION in ("chiRegularized", "tanhChi2", "chi2", "chi2xsec", "sigmoidChi2")
assert args.THEORYERROR >= 0
assert args.THEORYERROR <= 100.0
args.THEORYERROR = float(args.THEORYERROR)/100.0
if args.GOFFUNCTION != "chi2":
    print(f"WARNING: Support for GOFFUNCTION={args.GOFFUNCTION} is experimental and every output named chi2 or similar may not be a real chi2 but related to the GOFFUNCTION and may lose its meaning")
#  if args.GOFFUNCTION=="chi2xsec":
    #  print(f"WARNING: for GOFFUNCTION={args.GOFFUNCTION} only observable weights must be used, not individual bin or range weights")
#  # ## Get mandatory arguments
# if len(args) < 1:
#     args = "ipol.dat"
#    # print "Argument missing... exiting\n\n"
#    # op.print_usage()
#    # sys.exit(1)

if args.REFDIR is None:
    print("Error, no data directory specified (-d/--datadir), exiting\n\n")
    ap.print_usage()
    sys.exit(1)

if args.SCANSAMPLER not in ["uniform", "sobol", "latin"]:
    print("Error, requested sampling method '%s' not found, exiting" % args.SCANSAMPLER)
    import sys
    sys.exit(1)

if not os.path.exists(args.REFDIR):
    print("Error, specified data directory '%s' does not exist, exiting\n\n" % args.REFDIR)
    ap.print_usage()
    sys.exit(1)

if not os.path.exists(args.OUTDIR):
    os.makedirs(args.OUTDIR)


## Load Professor and show the standard banner
import professor2 as prof
if not args.QUIET:
    print(prof.logo)

# Read data files
if args.USE_RIVET_REF and args.WFILE:
    # TODO: support loading only histos for which there exist ipols,
    # but ipols are currently loaded later
    print("Reading reference histograms specified in weight file {args.WFILE}")
    if not os.path.exists(args.WFILE):
        print("Error: specified weights file '%s' does not exist, exiting\n\n" % args.WFILE)
        ap.print_usage()
        sys.exit(1)
    DHISTOS = prof.read_weightfile_histos(args.WFILE)
else:
    print(f"Reading all reference histograms from {args.REFDIR}...")
    import time
    t0 = time.time()
    DHISTOS = prof.read_all_histos(args.REFDIR)
    dt = time.time() - t0
    if dt > 10: #< longer than 10 seconds of loading
        print("Histogram loading is slow. Use a weight file for faster, targetted loading")


## Weight file parsing
if not args.WFILE:
    print("Warning: tuning with no weights file specified. Using all available bins with unit weight\n")
matchers = prof.read_pointmatchers(args.WFILE) if args.WFILE else None


## Try to read run histos and extract maximum errors --- NOTE this bit might be broken with patches NOTE
MAXERRDICT = None
if args.USE_RUNSDIR:
    try:
        _, RUNHISTOS = prof.read_all_rundata(args.RUNSDIR, None) #< don't care about reading params files
        MAXERRDICT = prof.find_maxerrs(RUNHISTOS)
    except:
        print("Could not read run data for error regularisation -- chi2 may be unstable")

## Slightly messy bit: load each ipol file we get, and store everything in a master dictionary
from collections import OrderedDict
MASTERBOX = OrderedDict()
MASTERCENTER = OrderedDict()
for ifile in args.IFILES:
    box, center, fitdata = prof.prepareBox(ifile, DHISTOS, matchers, MAXERRDICT, args.FILTER, args.DEBUG)
    if args.NORMW:
        totw = sum([x.w for x in fitdata[0][1]])
        for ibin in fitdata[0][1]:
            ibin.w /= totw
    MASTERBOX[box] = dict(fitdata)
    MASTERCENTER[center] = MASTERBOX[box]

## Take parameter names from the first box and assert that all other boxes (if present) have the same names
PNAMES = list(MASTERBOX.values())[0]["PNAMES"]
for v in list(MASTERBOX.values())[1:]:
    assert PNAMES == v["PNAMES"]

## Initial conditions --- use pos = center of hypercube, and step = range/10
pmins, pmaxs, pminmax = [], [], []
for num, pname in enumerate(PNAMES):
    testmin = [box[num][0] for box in MASTERBOX.keys()]
    testmax = [box[num][1] for box in MASTERBOX.keys()]
    pmins.append(min(testmin))
    pmaxs.append(max(testmax))
    pminmax.append([min(testmin), max(testmax)])
assert len(pmins) == len(pmaxs)

## Function definition wrapper
# TODO: allow user specified fitfunc
# TODO: use positional params and explicit names rather than signature detection
funcdef = prof.mkFitFunc("prof.simpleGoF", PNAMES, "profGoF",
                         ["MASTERBOX", "MASTERCENTER", "args.GOFFUNCTION", "args.THEORYERROR", "args.DEBUG"])
try:
    exec(funcdef, locals())
except SyntaxError as se:
    import textwrap
    print("Syntax error encountered in executing generated GoF function:\n")
    print(textwrap.indent(funcdef, "  ") + "\n")
    print("Check the generated variable names: are they valid Python? "
          "You should avoid use of 'raw' generator parameter names if "
          "they would include special characters, most prominently '.'"
          "or ':' are used by some event generators. You can use "
          "simpler variable names for tuning, and insert into the "
          "appropriate locations in templated scripts and steering cards")
    exit(20)
if args.DEBUG:
    print("Built GoF wrapper from:\n  '%s'" % funcdef)


## Unless we do a scan first, choose the center of param space as startpoint
pstart = [(pmins[i] + pmaxs[i])/2. for i in range(len(pmins))]

## Scan for better start point
# TODO: veto corner points with a ~hypersphere
if args.SCANNP:
    npoints = args.SCANNP
    startSampler = prof.NDSampler(pminmax, sampler=args.SCANSAMPLER, seed=args.SCANSEED)
    startPoints = [pstart] + [startSampler() for _ in range(npoints)] #< include the central point
    print("Scanning %i points" % npoints)
    testVals = [profGoF(*x) for x in startPoints]
    windex = testVals.index(min(testVals))
    winner = startPoints[windex]

    ## This sets the start point
    print("Using startpoint that yields fval %e:" % min(testVals))
    if windex == 0:
        print("Hypercube centre was the best start point: scan did not improve")
    for i, aname in enumerate(PNAMES):
        pstart[i] = winner[i]
        print("%s = %f" % (aname, pstart[i]))
    print()

## Dictionary fitarg for iminuit
import professor2 as prof
limits, fixed = prof.read_limitsandfixed(args.LIMITS)

# Set fixed parameters to their value
for i, aname in enumerate(PNAMES):
    if aname in list(fixed.keys()):
        pstart[i] = fixed[aname]
        print("Fix %s = %f" % (aname, pstart[i]))

if limits:
    print(f"Limiting parameters: {limits}")
if fixed:
    print(f"Fixing parameters: {fixed}")

FARG = prof.setupMinuitFitarg(PNAMES, pstart, pmins, pmaxs, limits, fixed, args.EXTRAPOLATE, verbose=not args.QUIET)

# print(FARG)
# print(profGoF)
# print(funcdef)
# if not args.QUIET:
#     print("\n" + 66*"*")
#     print("* Using iminuit, please visit https://github.com/iminuit/iminuit *")
#     print(66*"*" + "\n")

# TODO: errordef as CL params?
PRINTLEVEL = 0 if args.QUIET else 1

#minuit = Minuit(profGoF, errordef=1, print_level=PRINTLEVEL, forced_parameters=PNAMES, **FARG)
minuit = Minuit(profGoF, **FARG)
minuit.strategy = args.STRATEGY

## Set limits and fixed flags before launching minimization
if not args.EXTRAPOLATE and not args.QUIET:
    print(f"Limiting minimiser to interpolation limits for params {PNAMES}; you can use -x to allow extrapolation\n")
for i, pname in enumerate(PNAMES):
   if not args.EXTRAPOLATE:
       minuit.limits[pname] = (pmins[i],pmaxs[i])
   if pname in list(limits.keys()):
       minuit.limits[pname] = limits[pname]
   if pname in list(fixed.keys()):
       minuit.fixed[pname] = True

## Lift off
import time
start_time = time.time()
print("Tuning...")
minuit.migrad(args.NUM_CALLS, args.NUM_ATTEMPTS) #< max numbers of calls and attempts
if args.MINOS:
    minuit.minos()
else:
    minuit.hesse()
print("Minimisation finished after %2.1f seconds\n" % (time.time() - start_time))

## Now process the result:
## Goodness of fit
#print("%%%%", minuit.fixed)
n_fixedparams = len(fixed.keys())
#  n_freeparams = len(PNAMES) - len(minuit.list_of_fixed_param())
n_freeparams = len(PNAMES) - n_fixedparams
chi2 = minuit.fval
ndof = len(list(MASTERBOX.values())[0]["DBINS"]) - n_freeparams
ndof2 = 0
wdof2 = 0
for ibin in list(MASTERBOX.values())[0]["IBINS"]:
    if ibin.w > 0:
        ndof2 += 1
        wdof2 += ibin.w**2
ndof2 -= n_freeparams
wdof2 -= n_freeparams

# len(MASTERBOX.values()[0]["DBINS"]) - len(PNAMES) - len(minuit.list_of_fixed_param())
if not args.QUIET:
    print(f"GoF: {chi2:.2f}")
    print(f"Ndf_all: {ndof:d}, Ndf_w: {ndof2:d}, W2df: {wdof2:.2g}")
    if ndof and ndof2 and wdof2:
        print(f"GoF/Ndf_all: {chi2/ndof:.2g}, GoF/Ndf_w: {chi2/ndof2:.2g}, GoF/W2df: {chi2/wdof2:.2g}")

## Check if result is in validity range
result = [minuit.values[p] for p in PNAMES]
rok, rng = prof.is_inrange(result, pmins, pmaxs)
if not rok:
    msg = "Warning: parameters are outside the validity of the parametrization:"
    for i in rng:
        msg += "\n %s=%f ! in [%f,%f]" % (PNAMES[i], result[i], pmins[i], pmaxs[i])
    msg += "\n You might want to impose limits (--limits) on those parameters."
    if not args.QUIET:
        print(msg)

## Print weighted result
print("Best fit:")
print("\t".join(PNAMES))
print("\t".join(f"{r:.2g}" for r in result))

print()
print("Goodness of fit: %.2f" % chi2)

## Get the chi2 value setting all non-zero weights to 1
funcdefunit = prof.mk_fitfunc("prof.simpleGoF", PNAMES, "profGoFunit", ["MASTERBOX", "MASTERCENTER", "args.GOFFUNCTION", "args.THEORYERROR", "args.DEBUG", "True"])
# exec(funcdefunit, locals())
exec(funcdefunit, globals())
#print("Unit-weight GoF function function:", funcdefunit, "", sep="\n")
chi2unit = profGoFunit(*result)
print("Unweighted goodness of fit: %.2f" % chi2unit)
print()
print("Param values:", minuit.params, sep="\n")
print()
print("Param covariance matrix:", minuit.covariance, sep="\n")

## Max number of characters in any of parameter names --- for formatting (ljust)
LMAX = max([len(p) for p in PNAMES])

## Get the right box first
boxdict = prof.getBox(result, MASTERBOX, MASTERCENTER)

## Write out result
with open("%s/results.txt" % args.OUTDIR,"w") as f:
    ## Metadata
    f.write("# ProfVersion: %s\n" % prof.version())
    f.write("# Date: %s\n" % prof.mk_timestamp())
    if len(args.IFILES) == 1:
        f.write("# InterpolationFile: %s\n" % os.path.abspath(args.IFILES[0]))
    else:
        s = "# InterpolationFiles:"
        for ifile in args.IFILES:
            s += " %s" % os.path.abspath(ifile)
        f.write("%s\n" % s)
    f.write("# DataDirectory: %s\n" % os.path.abspath(args.REFDIR))

    ## Limits
    lstring = ""
    # pstate = minuit.get_param_states()
    pstate = minuit.params
    for num, p in enumerate(PNAMES):
        if pstate[num].has_limits and not minuit.fixed[p]:
            lstring += "\n#\t%s\t%f %f" % (p.ljust(LMAX), pstate[num].lower_limit, pstate[num].upper_limit)
    f.write("#\n# Limits:%s" % lstring)

    ## Fixed parameters
    fstring = ""
    for num, p in enumerate(PNAMES):
        if minuit.fixed[p]:
            fstring += "\n#\t%s\t%f" % (p.ljust(LMAX), pstate[num].value)
    f.write("\n#\n# Fixed:%s\n" % fstring)
    f.write("#\n# Minimisation result:\n#\n")
    f.write("# GOF %f\n" % chi2)
    f.write("# UNITGOF %f\n" % chi2unit)
    f.write("# NDOF %f\n" % ndof)

    ## Tuned parameter values
    for i, p in enumerate(PNAMES):
        f.write("%s\t%f\n" % (p.ljust(LMAX), minuit.values[PNAMES[i]]))
    with open(f"{args.OUTDIR}/tune.dat", "w") as tf:
        for i, p in enumerate(PNAMES):
            tf.write("%s\t%f\n" % (p.ljust(LMAX), minuit.values[PNAMES[i]]))
    np.save(f"{args.OUTDIR}/tune.npy", result)

    ## MIGRAD errors
    f.write("#\n# MIGRAD errors:\n#\n")
    for i, p in enumerate(PNAMES):
        f.write("# %s\t%e\n" % (p.ljust(LMAX), minuit.errors[PNAMES[i]]))

    ## MINOS stuff if we have it
    if minuit.merrors:
        f.write("#\n# MINOS errors:\n#\n")
        for i, p in enumerate(PNAMES):
            temp = sorted([k for k in list(minuit.merrors.keys()) if p == k[0]])
            minos = [minuit.merrors[k] for k in temp]
            #f.write("# %s\t%e\t%e\n" % (p.ljust(LMAX), minos[0], minos[1]))
            f.write("# %s\t%s\n" % (p.ljust(LMAX), minos))

    # Correlation matrix --- if params are fixed the covariance is not defined
    # The keys of minuit.covariance are tuples of strings
    f.write("#\n# Covariance matrix:\n#\n")
    try:
        import tabulate as tab
        f.write(tab.tabulate(*minuit.covariance.to_table(), tablefmt="plain") + "\n")
    except ImportError:
        for k, v in minuit.covariance.to_dict():
            f.write(f"{k}: {v}\n")
        f.write("\n")
        np.savetxt(f"{args.OUTDIR}/cov.dat", minuit.covariance)

    # t1, t2 = zip(*list(minuit.covariance))
    # l1 = list(t1)
    # CNAMES = list(set(l1))

    # from math import sqrt
    # for i in PNAMES:
    #     s="# %s" % i.ljust(LMAX)
    #     for j in PNAMES:
    #         if i in CNAMES and j in CNAMES:
    #             if minuit.covariance[(i,j)] >=0:
    #                 s+= "    %e"%(minuit.covariance[(i,j)]/sqrt(minuit.covariance[(i,i)]*minuit.covariance[(j,j)]))
    #             else:
    #                 s+= "   %e"%( minuit.covariance[(i,j)]/sqrt(minuit.covariance[(i,i)]*minuit.covariance[(j,j)]))
    #         else:
    #             s+= "    ---"
    #     f.write(s+"\n")
    # f.write("#\n# Covariance matrix:\n#\n")
    # for i in PNAMES:
    #     s="# %s"%i.ljust(LMAX)
    #     for j in PNAMES:
    #         if i in CNAMES and j in CNAMES:
    #             if minuit.covariance[(i,j)] >=0:
    #                 s+= "    %e"%minuit.covariance[(i,j)]
    #             else:
    #                 s+= "   %e"%minuit.covariance[(i,j)]
    #         else:
    #             s+= "    ---"
    #     f.write(s+"\n")


    # Weights ---  dump them all at the end
    f.write("#\n#\n# Weights used\n#\n")

    if matchers is None:
        for k in boxdict["IHISTOS"].keys():
            f.write("# %s\t1.0\n" % k)
    else:
        with open(args.WFILE) as g:
            for line in g:
                l = line.strip()
                if len(l) == 0 or l.startswith("#"):
                    continue
                f.write("# %s\n"%l)

    # Get the observables that actually entered the chi2
    f.write("#\n# Mean contribution:\n#\n")
    actual_obs = [k for k, v in boxdict["BINDICES"].items() if len(v)>0]
    for ao in actual_obs:
        tempf = dict(prof.boxFilt(boxdict, [ao]))
        tempb = { tempf["BOX"] : tempf}
        tempc = { tempf["CENTER"] : tempf}
        funcdefunit = prof.mk_fitfunc("prof.simpleGoF", PNAMES, "profGoFunit", ["tempb", "tempc", "args.GOFFUNCTION", "args.THEORYERROR", "args.DEBUG", "True"])
        exec(funcdefunit, locals())
        nbins_obs = len(boxdict["BINDICES"][ao])
        chi2obs = profGoFunit(*result)
        f.write("# %s\t%f\n" % (ao, chi2obs/nbins_obs))

# Minimiser profiles
if args.PROFILES:
    try:
        import pylab
        print("Drawing profiles")
        for num1, p1 in enumerate(PNAMES):
            pylab.clf()
            minuit.draw_profile(p1,subtract_min=True)
            pylab.savefig(os.path.join(args.OUTDIR, "profile_%s.pdf"%p1))
            pylab.clf()
            minuit.draw_mnprofile(p1,subtract_min=True)
            pylab.savefig(os.path.join(args.OUTDIR, "mnprofile_%s.pdf"%p1))
    except Exception as e:
        print("Could not draw profiles:")
        print(e)

# Minimiser contours
if args.CONTOURS:
    try:
        import pylab
        print("Drawing contours")
        pylab.clf()
        for num1, p1 in enumerate(PNAMES):
            for num2, p2 in enumerate(PNAMES):
                if num2 < num1:
                    pylab.clf()
                    try:
                        minuit.draw_contour(p1,p2)
                        pylab.savefig(os.path.join(args.OUTDIR, "contour_%s_%s.pdf"%(p1,p2)))
                    except:
                        print("Could not draw contour for %s vs. %s"%(p1,p2))
                        pass
                    pylab.clf()
                    try:
                        minuit.draw_mncontour(p1,p2)
                        pylab.savefig(os.path.join(args.OUTDIR, "mncontour_%s_%s.pdf" % (p1,p2)))
                    except:
                        print("Could not draw mncontour for %s vs. %s" % (p1,p2))
                        pass
    except Exception as e:
        print("Could not draw contours:")
        print(e)

## Write out ipolhistos
try:
    import yoda
    IHISTOS = boxdict["IHISTOS"]
    scatters = [IHISTOS[k].toDataHisto(result).toScatter2D() for k in sorted(IHISTOS.keys())]
    yoda.writeYODA(scatters, "{}/ipolhistos.yoda".format(args.OUTDIR))
except ImportError:
    print("Unable to import yoda, not writing out ipolhistos")

print()
print(f"Done. All output written to {args.OUTDIR}/")
