#! /usr/bin/python3

"""\
%(prog)s <runsdir> [datafile directory] [opts]

Make envelope plots showing the range coverage of the input run data.

TODO:
 - add an option to show translucent specific run-lines
"""

import argparse, sys
ap = argparse.ArgumentParser(usage=__doc__)
ap.add_argument("RUNSDIR", help="The directory containing the set of MC-run subdirs")
ap.add_argument("--pname", "--pfile", dest="PNAME", default="params.dat", help="Name of the params file to be found in each run directory (default: %(default)s)")
ap.add_argument("-w", "--wfile", dest="WFILE", default=None, help="Path to a weight file, used to restrict plotting to a subset of histograms (default: %(default)s)")
ap.add_argument("-o", "--outdir", dest="OUTDIR", default="envelopes", help="Output folder for plots (default: %(default)s)")
ap.add_argument("-d", "--datadir", "--refdir", dest="REFDIR", default=None, help="Data directory (default: %(default)s)")
ap.add_argument("-R", "--rivet", dest="USE_RIVET_REF", action="store_true", default=False, help="Use Rivet ref data (default: %(default)s)")
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")
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


## Load the Professor machinery
import professor2 as prof
if not args.QUIET:
    print((prof.logo))

## Load MC run histos and params
import os
if os.path.isdir(args.RUNSDIR):
    import glob
    indirs = glob.glob(os.path.join(args.RUNSDIR, "*"))
    try:
        PARAMS, HISTOS = prof.read_rundata(indirs, args.PNAME)
    except Exception as e:
        print(e)
        sys.exit(1)
else: # TODO: deprecate H5?
    try:
        PARAMS, HISTOS = prof.read_all_rundata_h5(args.RUNSDIR)
    except Exception as e:
        print(("Here wrong: {}".format(e)))
        sys.exit(1)

## Use a weight-file to select a histo subset
if args.WFILE:
    matchers = prof.read_pointmatchers(args.WFILE)
    for hn in list(HISTOS.keys()):
        if not any(m.match_path(hn) for m in list(matchers.keys())):
            del HISTOS[hn]
        elif args.DEBUG:
            print(("Observable {} passed weight file path filter".format(hn)))
    print(("Filtered observables by path, {} remaining".format(len(HISTOS))))
HNAMES = list(HISTOS.keys())

## If there's nothing left to plot, exit!
if not HNAMES:
    print("No observables remaining... exiting")
    sys.exit(1)


## Load plotting
import numpy as np
import matplotlib, os
matplotlib.use(os.environ.get("MPL_BACKEND", "Agg"))

def mk_envelope(histos):
    """ Take DataHistos and return coordinates for plotting """
    # Iterate over bins, get envelope representation
    # For each bin, return [xmin, xmax, mean(y), [percentiles...]]
    E = []
    for num, b in enumerate(histos[list(histos.keys())[0]].bins):
        t_b = np.array([x.bins[num].val for x in list(histos.values())])
        mean = np.nanmean(t_b)
        pctls = [np.nanpercentile(t_b, q) for q in (0, 5, 16, 84, 95, 100)]
        E.append([b.xmin, b.xmax, mean, pctls])  #min(t_b), max(t_b)])
    return E


def mk_data(histo):
    """ Take a DataHisto and return coordinates for plotting """
    d = []
    for num, b in enumerate(histo.bins):
        d.append([b.xmid, min(b.xmax-b.xmid, b.xmid-b.xmin), b.val, b.err])
    return d


def plot_envelope(env, data=None, name="envelope"):
    import matplotlib.pyplot as plt
    import matplotlib as mpl

    # This is the main figure object
    fig = plt.figure(figsize=(8,6), dpi=100)

    # This sets up the grid for main and ratio plot
    gs = mpl.gridspec.GridSpec(1, 1)#, height_ratios=[3,1], hspace=0)

    # Create a main plot
    axmain = fig.add_subplot(gs[0])
    axmain.set_ylabel("$f(x)$")
    axmain.set_xlabel("$x=$ %s" % name)

    # Convenience stuff for plotting clarity
    X, BAND_1SIG, BAND_2SIG, BAND_TOT, MEAN = [], [], [], [], []
    for num, b in enumerate(env):
        ymean = b[-2]
        ypcts  = b[-1]
        xmin  = b[0]
        xmax  = b[1]
        X.append(xmin)
        X.append(xmax)
        BAND_TOT.append([ypcts[0], ypcts[-1]])
        BAND_TOT.append([ypcts[0], ypcts[-1]])
        BAND_1SIG.append([ypcts[1], ypcts[-2]])
        BAND_1SIG.append([ypcts[1], ypcts[-2]])
        BAND_2SIG.append([ypcts[2], ypcts[-3]])
        BAND_2SIG.append([ypcts[2], ypcts[-3]])
        MEAN.append(ymean)
        MEAN.append(ymean)

    ## Use log y-axis if values are positive and there's a large dynamic range
    allvals = [b[0] for b in BAND_2SIG] + [b[1] for b in BAND_2SIG]
    if data is not None:
        allvals += [x[2] for x in data]
    allmin, allmax = min(allvals), max(allvals)
    with np.errstate(divide='ignore'):
        if np.log10(allmax)-np.log10(allmin) > 1.5:
            axmain.set_yscale("log")

    # Envelope
    axmain.fill_between(X, [b[0] for b in BAND_TOT],  [b[1] for b in BAND_TOT],  edgecolor="none", facecolor='yellow', alpha=0.5, interpolate=False)
    axmain.fill_between(X, [b[0] for b in BAND_2SIG], [b[1] for b in BAND_2SIG], edgecolor="none", facecolor='yellow', alpha=0.9, interpolate=False)
    axmain.fill_between(X, [b[0] for b in BAND_1SIG], [b[1] for b in BAND_1SIG], edgecolor="none", facecolor='gold', alpha=0.6, interpolate=False)
    # Mean of envelope
    axmain.plot(X, MEAN, color="orange", linewidth="2", linestyle="-", label="Mean")

    # Data plot
    if data is not None:
        Xdata = [x[0] for x in data]
        dX    = [x[1] for x in data]
        Ydata = [x[2] for x in data]
        dY    = [x[3] for x in data]
        axmain.errorbar(Xdata, Ydata, dY, dX, fmt="k.", linewidth=1.3, label="Data")

    # Switch off the frame in the legend
    leg = axmain.legend(loc=0, numpoints=1)
    fr = leg.get_frame()
    fr.set_visible(False)

    # Remove all unnecessary white space
    plt.tight_layout()

    # Output dir, path
    import os
    if not os.path.exists(args.OUTDIR):
        os.makedirs(args.OUTDIR)
    outname = os.path.join(args.OUTDIR, "%s.pdf" % name)

    # Save image as PDF
    plt.savefig(outname)
    plt.close(fig)


## Read reference-data histos
import os, sys, glob
DHISTOS = {}
if args.USE_RIVET_REF and args.WFILE:
    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)
elif args.REFDIR:
    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")

## Filter data histos to match the MC ones
DNAMES = {}
for h in HNAMES:
    for d in list(DHISTOS.keys()):
        if h in d:
            DNAMES[h] = d
## Free memory
for d in list(DHISTOS.keys()):
    if not d in list(DNAMES.values()):
        del DHISTOS[d]


# Plotting
for hname in list(HISTOS.keys()):
    h = HISTOS[hname]
    e = mk_envelope(h)
    # Sanitisation for label and output file name
    name = hname.lstrip("/").replace("/","_")

    # Plot ref-data too
    if args.REFDIR and hname in list(DNAMES.keys()):
        data = mk_data(DHISTOS[DNAMES[hname]])
        plot_envelope(e, data, name)
    else:
        plot_envelope(e, name=name)

if not args.QUIET:
    print(("Plots written to %s" % args.OUTDIR))
