#! /usr/bin/python3

"""
%(prog)s <param1,param2,param3> [<ipolfile>=ipol.dat]
or
%(prog)s PARAMFILE [<ipolfile>=ipol.dat]

Write out interpolated histograms at a given parameter point.

"""

import argparse, os, sys
ap = argparse.ArgumentParser(usage=__doc__)
ap.add_argument("PARAMSTR", help="Parameter string or filename")
ap.add_argument("IFILE", nargs="?", default="ipol.dat", help="Interpolation file (default: %(default)s)")
ap.add_argument("-o", "--output", dest="OUTPUT", default="pred.yoda", help="Output file name (default: %(default)s)")
ap.add_argument("--wfile", dest="WFILE", default=None, help="Path to a weight file, used to restrict output to a subset of histograms (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()


## Imports
try:
    import yoda
except ImportError:
    "YODA not found, exiting"
    import sys
    sys.exit(1)
import professor2 as prof
if not args.QUIET:
    print(prof.logo)


## Parse the param-point argument
PARAMS = None
if os.path.exists(args.PARAMSTR):
    NAMES, VALS = [], []
    with open(args.PARAMSTR) as f:
        for l in f:
            lparts = l.strip().split()
            if l.startswith("#") or not lparts: continue
            NAMES.append( lparts[0] )
            VALS.append( float(lparts[-1]) )
else:
    print("No parameter file given or specified param file does not exist, exiting...")
    sys.exit(1)

PARAMS = dict(zip(NAMES, VALS))

## Read ipol param names and histos
IHISTOS, META = prof.read_ipoldata(args.IFILE)
PNAMES = META["ParamNames"].split()

# Reorder parampoint values according to order of params in ipol file
TESTPOINT = [PARAMS[k] for k in PNAMES]

## Weight file parsing
if args.WFILE:
    matchers = prof.read_pointmatchers(args.WFILE)
    for hn in list(IHISTOS.keys()):
        if not any(m.match_path(hn) for m in list(matchers.keys())):
            del IHISTOS[hn]
    if len(list(IHISTOS.keys()))==0:
        print("Nothing left after weight file parsing, exiting")
        sys.exit(0)


## Write out ipolhistos
ofile = args.OUTPUT
scatters=[IHISTOS[k].toDataHisto(TESTPOINT).toScatter2D() for k in sorted(IHISTOS.keys())]
yoda.writeYODA(scatters, ofile)

if not args.QUIET:
    print("Wrote output to", ofile)
