#! /usr/bin/python3

"""\
%(prog)s <runsdir> [<ipolfile>=ipol.dat] [opts]

Interpolate histo bin values as a function of the parameter space by loading
run data and parameter lists from run directories in $runsdir (often "mc")

TODO:
 * Use weight file position matches to exclude some bins, as well as path matching
 * Handle run combination file/string (write a hash of the run list into the ipol filename?)
 * Support asymm error parametrisation
"""

import argparse, os, sys
ap = argparse.ArgumentParser(usage=__doc__)
ap.add_argument("RUNSDIR", help="Path to the runs directory to parametrize")
ap.add_argument("IFILE", nargs="?", default="ipol.dat", help="Path to the output interpolation file (default: %(default)s)")
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("--limits", dest="LIMITS", default=None, help="Simple text file with parameter limits and fixed parameters")
ap.add_argument("-w", "--wfile", dest="WFILE", default=None, help="Path to a weight file, used to restrict ipol building to a subset of bins (default: %(default)s)")
ap.add_argument("--order", dest="ORDER", default=3, help="Global order of polynomials for interpolation. auto also possible")
ap.add_argument("--ierr",  dest="ERR_MODE", default="median", help="How to include MC errors: none, mean, median, symm (parametrised) (default: %(default)s)") #< add rel, asymm
ap.add_argument("--eorder", dest="ERR_ORDER", default=None, type=int, help="Global order of polynomials for uncertainty interpolation (default: same as from --order)")
ap.add_argument("--rc", dest="RUNCOMBS", default=None, help="Run-combination file")
# TODO: Add a no-scale option
# TODO: Change to instead (optionally) specify the max number of parallel threads/procs. Use by default if ncores > 1?
ap.add_argument("-j", dest="MULTI", type=int, default=1, help="Number of threads to use")
ap.add_argument("--summ",  dest="SUMMARY", default="", help="Summary description to be written to the ipol output file")
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("--split", dest="SPLIT", default=None, help="Can incorporate a linear split in parameter space. Provide \"ParamName1 ParamName2 Value1 Value2 Gradient (these describe the line down which to split the plot. Value1 and Value2 form the coords of a point on the line.) ABOVE/BELOW (ABOVE => use runs above the line etc)\"")
ap.add_argument("--subspace", dest="SUBSPACE", default=None, help="Simple x<y filtering. Provide \"ParamName ABOVE/BELOW Value\"")
ap.add_argument("--logy", dest="LOGY", action="store_true", default=False, help="Parametrize the logarithm of the read-in values. (Experimental)")
ap.add_argument("--auto-omin",  dest="AUTO_OMIN", type=int, default=0,  help="Minimum allowed order with --order auto")
ap.add_argument("--auto-omax",  dest="AUTO_OMAX", type=int, default=4,  help="Maximum allowed order with --order auto")
ap.add_argument("--auto-nit",   dest="AUTO_NIT",  type=int, default=10, help="Number of iteration with --order auto")
ap.add_argument("--auto-split",   dest="AUTO_SPLIT",  type=float, default=0.2, help="Fraction of test sample with --order auto")
ap.add_argument("--medianfilt",   dest="MEDIAN_FILT",  type=float, default=-1, help="Bin-wise filtering of inputs by median relative error. (Experimental)")
args = ap.parse_args()

if not os.path.exists(args.RUNSDIR):
    print("Error, specified run dir '%s' does not exist" % args.RUNSDIR)
    sys.exit(1)

## By default, use the same ipol order for errors as for values
if args.ERR_ORDER is None:
    args.ERR_ORDER = args.ORDER

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

# Read details of split in parameter space if necessary
if args.SPLIT != None:
    Split = map(str, args.SPLIT.split(' '))
    splitParam1 = Split[0]
    splitParam2 = Split[1]
    splitVal1 = float(Split[2])
    splitVal2 = float(Split[3])
    splitGrad = float(Split[4])

## Load MC run histos and params
try:
    if args.RUNSDIR.endswith(".yaml"):
        try:
            import yaml
        except ImportError as e:
            print(("Unable to import YAML:", e))
            import sys
            sys.exit(1)
        PARAMS, HISTOS = prof.read_all_rundata_yaml(args.RUNSDIR)
    elif args.RUNSDIR.endswith(".h5") or args.RUNSDIR.endswith(".hdf5"):
        try:
            import h5py
        except ImportError as e:
            print(("Unable to import H5PY:", e))
            import sys
            sys.exit(1)
        PARAMS, HISTOS = prof.read_all_rundata_h5(args.RUNSDIR)
    else:
        PARAMS, HISTOS = prof.read_all_rundata(args.RUNSDIR, args.PNAME, ignore_missing_histos=True)
    RUNS, PARAMNAMES, PARAMSLIST = prof.mk_ipolinputs(PARAMS)
except Exception as e:
    print(e)
    sys.exit(1)

# --rc runcombs.dat:4
# would use the 4th line of the runcombs.dat file
if args.RUNCOMBS is not None:
    f_rc, line = args.RUNCOMBS.split(":")
    with open(f_rc) as f:
        temp=[l for l in f]
    thisRC = temp[int(line)].split()

    # Filtering and overwriting
    thisRUNS, thisPARAMSLIST = [], []
    for num, r in enumerate(RUNS):
        if r in thisRC:
            thisRUNS.append(r)
            thisPARAMSLIST.append(PARAMSLIST[num])
    RUNS=thisRUNS
    PARAMSLIST=thisPARAMSLIST


## Some useful announcements about the data loaded and the interpolation planned
if not args.QUIET and not args.ORDER=="auto":
    thisO = int(args.ORDER)
    minnruns =prof.calcnumCoeffs(len(PARAMNAMES), thisO)
    if (len(RUNS) < minnruns):
        print("Not enough runs for order %i polynomials --- would require %i" % \
              (thisO, minnruns))
        for i in range(1, thisO):
            if (prof.calcnumCoeffs(len(PARAMNAMES), thisO -i) <= len(RUNS)):
                print("Try order %i (min %i runs)" % \
                      (thisO -i, prof.calcnumCoeffs(len(PARAMNAMES), thisO -i)))
        import sys
        sys.exit(1)
    else:
        print("Building order-%d interpolations in %d params: require at least %d runs" % \
            (thisO, len(PARAMNAMES), prof.calcnumCoeffs(len(PARAMNAMES), thisO)))
        print("Loaded %d distinct observables from %d runs" % (len(HISTOS), len(RUNS)))


## Weight file parsing to select a histos 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 %s passed weight file path filter" % hn)
    print("Filtered observables by path, %d remaining" % len(HISTOS))
HNAMES = list(HISTOS.keys())





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


## Robustness tests and cleaning: only retain runs that contain every histo
# TODO: combine with weights histo vetoing -- should we explicitly unload unused run data, or keep it for the next combination to use? Or do we now leave runcombs to the user?
bad, badnum = [], []
for irun, run in enumerate(RUNS):
    for hn in HNAMES:
        if not run in  HISTOS[hn]:#.has_key(run):
            bad.append(run)
            badnum.append(irun)
            break
if args.LIMITS != None:
    limits, fixed = prof.read_limitsandfixed(args.LIMITS)
    for irun, run in enumerate(RUNS):
        for inam, nam in enumerate(PARAMNAMES):
            ## Ignore outside-limits detection if parameter is fixed
            if nam not in limits.keys():
                continue
            if PARAMSLIST[irun][inam] <= limits[nam][0] or PARAMSLIST[irun][inam] >= limits[nam][1]:
                bad.append(run)
                badnum.append(irun)
                break

# if args.SPLIT != None:
    # ip1 = PARAMNAMES.index(splitParam1)
    # ip2 = PARAMNAMES.index(splitParam2)
    # for irun, run in enumerate(RUNS):
        # ycheck = PARAMSLIST[irun][ip2]-splitGrad*(PARAMSLIST[irun][ip1]-splitVal1)
	# if Split[5] == "ABOVE": #if taking values above the split, remove those below
        # if ycheck <= splitVal2:
            # bad.append(run)
            # badnum.append(irun)
	# elif Split[5] == "BELOW": #if taking values below the split, remove those above
        # if ycheck >= splitVal2:
            # bad.append(run)
            # badnum.append(irun)
	# else:
        # print("Error: must specify either ABOVE or BELOW the split")
        # sys.exit(1)

# if args.SUBSPACE is not None: #NOTE for future, do this with text files
    # pname, operation, value = args.SUBSPACE.split()
    # ip = PARAMNAMES.index(pname)
    # value=float(value)
    # for irun, run in enumerate(RUNS):
        # ycheck = PARAMSLIST[irun][ip]
	# if operation == "ABOVE": #if taking values above the split, remove those below
        # if ycheck <= value:
            # bad.append(run)
            # badnum.append(irun)
	# elif operation == "BELOW": #if taking values below the split, remove those above
        # if ycheck >= value:
            # bad.append(run)
            # badnum.append(irun)
	# else:
        # print("Error: must specify either ABOVE or BELOW the split")
        # sys.exit(1)


if bad:
    print("Found %d bad runs in %d total... removing" % (len(bad), len(RUNS)))
    goodr, goodp = [], []
    for irun, run in enumerate(RUNS):
        if not irun in badnum:
            goodr.append(run)
            goodp.append(PARAMSLIST[irun])
    RUNS = goodr
    PARAMSLIST = goodp

## If there's nothing left to interpolate, exit!
if not RUNS:
    print("No valid runs remaining... exiting")
    sys.exit(1)

CONFIG = {"MULTI":args.MULTI, "ORDER":args.ORDER, "ERR_MODE":args.ERR_MODE, "ERR_ORDER":args.ERR_ORDER, "QUIET":args.QUIET, "DEBUG":args.DEBUG}
CONFIG["AUTO_OMIN"] = args.AUTO_OMIN
CONFIG["AUTO_OMAX"] = args.AUTO_OMAX
CONFIG["AUTO_NIT"] = args.AUTO_NIT
CONFIG["AUTO_SPLIT"] = args.AUTO_SPLIT
CONFIG["MEDIAN_FILT"] = args.MEDIAN_FILT
CONFIG["LOGY"]=args.LOGY
IPOLDICT = prof.mkStandardIpols(HISTOS, HNAMES, RUNS, PARAMSLIST, CONFIG)

# from IPython import embed
# embed()

## Active memory clean-up
del HISTOS

## Write out meta info
prof.writeIpol(args.IFILE, IPOLDICT, [PARAMNAMES, PARAMSLIST], RUNS, args.SUMMARY, args.RUNSDIR)
