#! /usr/bin/python3

"""\
%prog <refdir> [<ipolfile>=ipol.dat] [<runsdir>=<refdir>/../mc] [opts]

This is an example script that uses gp_minimize from scikit-optimize
as minimiser which is not gradient based.

This is experimental work.
"""

import optparse, os, sys
op = optparse.OptionParser(usage=__doc__)
op.add_option("-d", "--datadir", dest="DATADIR", default=None, help="The data directory")
op.add_option("-o", "--outdir", dest="OUTDIR", default="tunes", help="Prefix for outputs (default: %default)")
op.add_option("--wfile", dest="WFILE", default=None, help="Path to a weight file to specify unequal chi2 weights of each bin in the fit (default: %default)")
op.add_option("--limits", dest="LIMITS", default=None, help="Simple text file with parameter limits and fixed parameters")
op.add_option("--gp-np", dest="NPOINTS", default=100, type=int, help="Number of evaluations")
op.add_option("-v", "--debug", dest="DEBUG", action="store_true", default=False, help="Turn on some debug messages")
op.add_option("-q", "--quiet", dest="QUIET", action="store_true", default=False, help="Turn off messages")
op.add_option("--filter", dest="FILTER", action="store_true", default=False, help="Filter out data bins that have 0 error")
op.add_option("-r", "--runsdir", dest="RUNSDIR", default=None, help="The runs directory")
op.add_option("--limit-errs", dest="USE_RUNSDIR", action="store_true", default=False, help="Re-read the runsdir to regularise error ipols")
opts, args = op.parse_args()


try:
    from skopt import gp_minimize
except ImportError:
    print "Cannot find skopt, try pip install scikit-optimize"

## Get mandatory arguments
if len(args) < 1:
    print "Argument missing... exiting\n\n"
    op.print_usage()
    sys.exit(1)
REFDIR = opts.DATADIR
if REFDIR is None:
    print "Error, no data directory specified (-d/--datadir), exiting\n\n"
    op.print_usage()
    sys.exit(1)
IFILES = args
RUNSDIR = opts.RUNSDIR

if not os.path.exists(opts.OUTDIR):
    try:
        os.makedirs(opts.OUTDIR)
    except:
        pass

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

# Read data files
DHISTOS = prof.read_all_histos(REFDIR)

## Weight file parsing
matchers = prof.read_pointmatchers(opts.WFILE) if opts.WFILE else None


## Try to read run histos and extract maximum errors --- NOTE this bit might be broken with patches NOTE
MAXERRDICT = None
if opts.USE_RUNSDIR:
    try:
        _, RUNHISTOS = prof.read_all_rundata(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 IFILES:
    box, center, fitdata = prof.prepareBox(IFILE, DHISTOS, matchers, MAXERRDICT, opts.FILTER, opts.DEBUG)
    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=MASTERBOX.values()[0]["PNAMES"]
for v in MASTERBOX.values()[1:]:
    assert PNAMES == v["PNAMES"]


## Function definition wrapper
funcdef = prof.mk_fitfunc("prof.simpleGoF", PNAMES, "profGoF", ["MASTERBOX", "MASTERCENTER", "opts.DEBUG"])
exec funcdef in locals()
if opts.DEBUG:
    print "Built GoF wrapper from:\n  '%s'" % funcdef

if not opts.QUIET:
    print "\n"
    print 90*"*"
    print "* Using scikit-optimize, please visit https://github.com/scikit-optimize/scikit-optimize *"
    print 90*"*"
    print "\n"


## Determine the box for the sampler

pmins, pmaxs = [], []
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))
assert len(pmins) == len(pmaxs)


## Fix parameters, set limits (with pname translation)
limits, fixed = prof.read_limitsandfixed(opts.LIMITS)
if len(limits)>0:
    for num, pname in enumerate(PNAMES):
        if pname in limits.keys():
            pmins[num] = limits[pname][0]
            pmaxs[num] = limits[pname][1]

funcdef = prof.mk_fitfunc("prof.simpleGoF", PNAMES, "profGoF", ["MASTERBOX", "MASTERCENTER", "opts.DEBUG"])
exec funcdef in locals()

import time
start_time = time.time()

from skopt import gp_minimize
res = gp_minimize(lambda x: prof.simpleGoF(x, MASTERBOX, MASTERCENTER, opts.DEBUG), zip(pmins, pmaxs), n_calls=opts.NPOINTS)
print("\ngp_minimize finished after %s seconds\n" % (time.time() - start_time))

# Some output
print("Minimum value found: %e \nat:\n"%res['fun'])
for num, pname in enumerate(PNAMES):
    print pname, res['x'][num]



with open(os.path.join(opts.OUTDIR, "gpresult.txt"), "w") as f:
    f.write("# Minimum: %e\n"%(-res['fun']))
    for x in zip(PNAMES,res['x']):
        f.write("%s\t%e\n"%(x[0], x[1]))

try:
    import yoda
    result = res["x"]
    boxdict = prof.getBox(result, MASTERBOX, MASTERCENTER)
    IHISTOS=boxdict["IHISTOS"]
    scatters=[IHISTOS[k].toDataHisto(result).toScatter2D() for k in sorted(IHISTOS.keys())]
    yoda.writeYODA(scatters, "%s/gp_ipolhistos.yoda" % opts.OUTDIR)
except ImportError:
    print "Unable to import yoda, not writing out ipolhistos"

print "Output written to", opts.OUTDIR
