#! /usr/bin/python3

"""
%(prog)s [-t template1.txt -t ...] [-o out1] <myrangefile|--extend>

Sample a parameter space, creating a set of parameter files and optionally
substituting into script templates, with either uniform sampling (default) or
sampling in a transformed space.

Parameter ranges (and bias functions) can be given in a parameter-range file
with one parameter per line and whitespace separation of the name, low, high,
and bias terms. Alternatively, the 'extend' mode allows adding of new template
instantiations to existing scan/run directories.
"""

import argparse
ap = argparse.ArgumentParser(usage=__doc__,formatter_class=argparse.ArgumentDefaultsHelpFormatter)
ap.add_argument("-n", dest="NUMPOINTS", metavar="NUM", type=int, default=100, help="number of samples to generate")
ap.add_argument("-t", dest="TEMPLATES", metavar="FILE", action="append", default=[], help="specify a template file to be populated for each sample point. Can be given multiple times. Param names and the run-number N, given in curly braces, are expanded to the sampled values using the Python format mini-language.")
ap.add_argument("-o", "--outdir", dest="OUTDIR", metavar="DIR", default="scan", help="specify the output directory name ")
ap.add_argument("-p", "--pfile", dest="PFNAME", metavar="FILE", default="params.dat", help="specify params file base name to be populated for each sample point")
#ap.add_argument("-m", "--outmode", dest="OUTMODE", metavar="MODE", default="hier", help="specify the output structuring mode: either 'hier' (default) or 'table' to create one text file with a table like structure")
ap.add_argument("--sampler", dest="SAMPLER", default="uniform", help="Select sampling method uniform|latin|sobol|grid|normal|extend ")
ap.add_argument("--normal-spec", dest="NORMAL_SPEC", metavar="FILE", default=None, help="specify comma-separated mean-vector file, covariance-matrix file, and an std scaling factor for 'normal' sampling mode")
ap.add_argument("--seed", dest="SEED", metavar="VAL", type=int, default=None, help="Random seed for the sampler ")
ap.add_argument("--overlap-fraction", dest="OVERLAP", default=0, type=float, help="Fraction (0.1 means 10 per cent) of axis overlap when sampling multiple patches ")
ap.add_argument("--extend", dest="EXTEND", action="store_true", default=False, help="Select the 'extend' sampling method to add more templates to existing scan dirs")
ap.add_argument("--veto", dest="VETOFN", metavar="FILE", default=None, help="specify a file from which to read the definition of a Python sample point-vetoing function, prof_sample_veto(paramsdict). Return True means 'veto point'")
ap.add_argument("--hsphere", dest="HSPHERE", action="store_true", default=False, help="Veto samplings outside a hypersphere (in normalized params) so high-dim samplings are not dominated by the hypercube 'hypercorners'")
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("RANGEFILE", nargs="?", default=None, help="a parameter-range file with <name> <low> <high> [<bias>] per line")
args = ap.parse_args()

import sys, os
#assert args.OUTMODE in ("hier", "table")
if not args and not args.EXTEND and not args.SAMPLER == "extend":
    print("Error: no arguments given\n")
    ap.print_help()
    sys.exit(1)

## Handle args relating to non-default samplers
if args.EXTEND:
    args.SAMPLER = "extend"
if args.SAMPLER not in ["uniform", "sobol", "latin", "grid", "extend", "normal"]:
    print("Error, requested sampling method '%s' not found, exiting" % args.SAMPLER)
    sys.exit(1)
## Modify args for Sobol sampling
if args.SAMPLER == "sobol":
    try:
        import sobol
    except (ImportError, ModuleNotFoundError) as e:
        print("Extra modules are required to perform quasirandom sampling: try 'pip install sobol sobol_seq'")
        exit(1)
    if args.SEED and args.SEED > 100000:
        print("Warning, Sobol sampling is slow for seed > 100000, current seed: %i" % args.SEED)
    import math
    log2npts = math.log2(args.NUMPOINTS)
    if not log2npts.is_integer():
        newnpts = 2**math.ceil(log2npts)
        print(f"Warning, Sobol sampling requires 2**N samples for statistical balance: upgrading from {args.NUMPOINTS:d} to {newnpts:d} samples")
        args.NUMPOINTS = newnpts
## Create a multivariate normal sampler
if args.SAMPLER == "normal":
    assert args.NORMAL_SPEC is not None
    meanfile, covfile, scale = args.NORMAL_SPEC.split(",")
    mean = np.loadtxt(meanfile)
    cov = np.loadtxt(covfile)
    cov *= float(scale)**2
    from scipy.stats import multivariate_normal as mvnorm
    NORMAL_SAMPLER = mvnorm(mean, cov, seed=args.SEED)

import numpy as np
try:
    from collections import OrderedDict
except:
    from ordereddict import OrderedDict

def mkrunstr(num):
    return "{run:04d}".format(run=num)

def mkoutname(fname, run, prefix="", suffix=""):
    if type(run) is int:
        run = mkrunstr(run)
    # if args.OUTMODE == "hier":
    name = os.path.join(run, fname)
    # elif args.OUTMODE == "flat":
    #     fname = os.path.basename(fname)
    #     base, ext = os.path.splitext(fname)
    #     name = prefix + base + "-" + run + suffix + ext
    #     print("Name:", name)
    # elif args.OUTMODE == "table":
    #     name = fname
    return os.path.join(args.OUTDIR, name)

def mkdir(path):
    d = os.path.dirname(path) #< if path is to a dir, make sure to terminate with a /
    if not os.path.exists(d):
        os.makedirs(d)

def sample(sdict, N, vetofns, sampler="uniform", seed=None):
    ## Do random param sampling and template instantiation
    ppoints = []
    ranges = [[float(x[0].low), float(x[0].high)] for x in list(sdict.values())]
    biases = [x[1] for x in list(sdict.values())]
    pnames = list(sdict.keys())

    ## Make a sampler, either internal or an augmented SciPy sampler (probably MV Gaussian)
    if type(sampler) is str:
        from professor2 import NDSampler
        nds = NDSampler(ranges, biases, sampler=sampler, seed=seed)
    elif hasattr(sampler, "rvs"):
        nds = sampler
        #nds.__call__ = lambda self : return self.rvs(size=1)

    ## Call the sampler to produce N valid points
    for n in range(N):
        # npad = mkrunstr(n)

        ## Populate params dictionaries
        while True: #< until the veto is passed
            params = OrderedDict.fromkeys(pnames)
            if hasattr(nds, "rvs"):
                point = nds.rvs()
            else:
                point = nds()
            for num, p in enumerate(pnames):
                params[p] = point[num]

            ## Allow user functions to veto the point
            vetoed = False
            for vf in vetofns:
                if vf(params):
                    vetoed = True
                    break
            if vetoed:
                continue #< try sampling again

            ## Successful sample: record and continue
            ppoints.append(params)
            break

    return ppoints


# https://stackoverflow.com/questions/12864445/numpy-meshgrid-points/12891609
def meshgrid2(*arrs):
    import numpy as np
    arrs = tuple(reversed(arrs))
    lens = map(len, arrs)
    dim = len(arrs)
    sz = 1
    for s in lens:
       sz *= s
    ans = []
    for i, arr in enumerate(arrs):
        slc = [1]*dim
        slc[i] = lens[i]
        arr2 = np.asarray(arr).reshape(slc)
        for j, sz in enumerate(lens):
            if j != i:
                arr2 = arr2.repeat(sz, axis=j)
        ans.append(arr2)
    return tuple(ans)


def grid(sdict, N):
    """
    Create a grid of points in the paramspace.
    """
    ranges = [[float(x[0].low), float(x[0].high)] for x in list(sdict.values())]
    pnames = list(sdict.keys())
    if type(N) is list:
        assert(len(N) == len(pnames))

    lsps = OrderedDict()
    for num, pname in enumerate(pnames):
        lsps[pname] = np.linspace(*ranges[num], num=N)
    #print(lsps)

    G = meshgrid2(*list(lsps.values()))
    Gpoints = np.vstack(map(np.ravel, G)).T
    ppoints = []
    for v in Gpoints:
        od = OrderedDict.fromkeys(pnames)
        for num, pname in enumerate(pnames):
            od[pname] = v[num]
        ppoints.append(od)
    return ppoints


def writeParams(p, fpath):
    with open(fpath, "w") as pf:
        for k, v in list(p.items()):
            pf.write("{name} {val:e}\n".format(name=k, val=v))


def writeTemplates(p, rundir, templates):
    for tbasename, (tmpl, mode) in list(templates.items()):
        ## Perform substitution using standard Python formatting mini-lang
        ## or if this fails try to parse the template with sympy if availible
        try:
            txt = tmpl.format(**p)
        except KeyError:
            try:
                from sympy import parse_expr
            except (ModuleNotFoundError, ImportError) as e:
                print(e)
                print("Please install sympy using 'pip install --user sympy'")
                sys.exit(1)
            txt = tmpl
            for lin in txt.split("\n"):
                if -1 == lin.find("{"):
                    continue
                for i in list(p.keys())[:-1]:
                    lin = lin.replace(i,str(p[i]))
                    txt = txt.replace(i,str(p[i]))
                expr = lin.split('{')[1].split('}')[0]
                res = parse_expr(expr)
                txt = txt.replace("{"+expr+"}",str(res))
        ## Write out
        tpath = os.path.join(rundir, tbasename)
        if args.DEBUG:
            print(f"Writing {tpath} template instantiation")
        with open(tpath, "w") as tf:
            tf.write(txt)
        os.chmod(tpath, mode)


def readParams(scandir, pfname):
    "Read the point number and params from each subdir, and return the param dict for each"

    with os.scandir(scandir) as odentries:
        ppoints = {}
        for entry in odentries:
            if entry.is_dir():
                pnum = os.path.basename(entry.name)
                ppath = os.path.join(scandir, pnum, pfname)
                #print(entry.name, os.path.basename(entry.name), ppath, os.path.exists(ppath))
                if not os.path.exists(ppath):
                    continue

                params = OrderedDict()
                try:
                    with open(ppath) as pfile:
                        for line in pfile:
                            pname, pval = line.strip().split(maxsplit=1)
                            pval = float(pval)
                            params[pname] = pval
                        params["N"] = pnum #< run num is not stored in the
                        ppoints[pnum] = params
                except Exception as e:
                    print(f"Issue in extend-mode reading params from {ppath}... skipping this point")
                    continue
        ppoints = [v for k,v in sorted(ppoints.items())]
        return ppoints


## Populate dict of script templates
templates = {}
for tpath in args.TEMPLATES:
    ## Remove a .tmpl extension from the instantiations
    tname = os.path.basename(tpath)
    if tname.endswith(".tmpl"):
        tname = tname[:-5]
    ## Get the template-file mode, to be copied to instantiations
    import stat
    mode = stat.S_IMODE(os.stat(tpath).st_mode)
    ## Read the template and store
    with open(tpath, "r") as f:
        templates[tname] = (f.read(), mode)


## Populate param samplers dict
samplers = OrderedDict()
if args.SAMPLER != "extend":
    from professor2 import Sampler
    with open(args.RANGEFILE, "r") as prf:
        for line in prf:

            ## Read the sampling spec
            line = line.split("#")[0].strip() #< strip comments and whitespace
            if line:
                parts = line.split()
                name = parts[0]
                samplers[name] = []
                try:
                    float(parts[-1])
                    bias = None
                except ValueError:
                    bias = parts[-1]

                ## Make the (sub)samplers
                Nsub = len(parts)-1 if bias is None else len(parts)-2
                for i in range(Nsub-1):
                    if len(parts[1:]) > 2:
                        distance = float(parts[2+i]) - float(parts[1+i])
                        dx = args.OVERLAP * distance
                        ## Left edge stays left edge
                        if i == 0:
                            subsampler = Sampler(float(parts[1+i]), float(parts[2+i])+dx)
                        ## Right edge stays right edge
                        elif i == len(parts[1:])-2:
                            subsampler = Sampler(float(parts[1+i])-dx, float(parts[2+i]))
                        ## Some guy in the middle, add overlap on both sides
                        else:
                            subsampler = Sampler(float(parts[1+i])-dx, float(parts[2+i])+dx)
                    else: ## For parameter axes with only one patch...
                        subsampler = Sampler(float(parts[1+i]), float(parts[2+i]))

                    ## Add to the list
                    if args.DEBUG:
                        print(name, subsampler)
                    samplers[name].append( (subsampler, bias) )


## Use only one dummy patch dict for the extend mode
if args.SAMPLER == "extend":
    #print(samplers, patches, patchDicts)
    patchDicts = [None]
else:
    ## Otherwise make all the patches from range combinations
    # https://stackoverflow.com/questions/798854/all-combinations-of-a-list-of-lists
    import itertools
    patches = list(itertools.product(*list(samplers.values())))
    ## For each patch, make a dictionary that the NDSampler understands
    patchDicts = []
    for patch in patches:
        stemp = OrderedDict.fromkeys(list(samplers.keys()))
        for nparam, param in enumerate(stemp.keys()):
            stemp[param] = patch[nparam]
        patchDicts.append(stemp)


VETOFNS = []

## Load a point veto function if supplied by the user
if args.VETOFN:
    # Python2
    # execfile(args.VETOFN)
    # Suggested replacement by https://stackoverflow.com/questions/6357361/alternative-to-execfile-in-python-3
    with open(args.VETOFN, "rb") as source_file:
        code = compile(source_file.read(), args.VETOFN, "exec")
    exec(code)
    assert "prof_sample_veto" in dir()
    VETOFNS.append(prof_sample_veto)

## Add a (unbiased) hypersphere (ellipse) veto
if args.HSPHERE:
    PMINMAX = {}
    for pname in samplers.keys():
        for subsampler, bias in samplers[pname]:
            #print(subsampler)
            PMINMAX.setdefault(pname, [subsampler.low, subsampler.high])
            PMINMAX[pname][0] = min(subsampler.low, PMINMAX[pname][0])
            PMINMAX[pname][1] = max(subsampler.high, PMINMAX[pname][1])
    # TODO: is there a way to force capture by value? Is it needed?
    # TODO: correct for bias??
    def prof_hsphere_veto(ppt):
        normpvals = []
        for pname, pval in ppt.items():
            pmm = PMINMAX[pname]
            pmid = (pmm[1] + pmm[0]) / 2.0
            pscale = (pmm[1] - pmm[0]) / 2.0
            normpval = (pval-pmid) / pscale
            normpvals.append(normpval)
        if sum(npv**2 for npv in normpvals) > 1:
            return True
    VETOFNS.append(prof_hsphere_veto)


## Use a time-based seed by default (not good for repetition/extension)
if args.SEED:
    useSEED = args.SEED
else: # if 0 or unset
    import time
    useSEED = int(time.time())
    if args.SAMPLER == "sobol":
        useSEED = useSEED % 100000
        print("INFO: sobol doesn't like large seeds, converted system time seed to %i" % useSEED)


## Sample points and write/print out
if args.SAMPLER == "extend":
    print(f"Extending content of sampled points in {args.OUTDIR}/")
else:
    print(f"Writing sampled point data to {args.OUTDIR}")
table = ""
for npatch, patch_sdict in enumerate(patchDicts):
    if args.SAMPLER == "extend":
        #assert args.OUTMODE == "hier"
        pps = readParams(args.OUTDIR, args.PFNAME)
    elif args.SAMPLER == "grid":
        # Note: no veto, etc.
        pps = grid(patch_sdict, args.NUMPOINTS)
    elif args.SAMPLER == "normal":
        pps = sample(patch_sdict, args.NUMPOINTS, VETOFNS, NORMAL_SAMPLER)
    else:
        pps = sample(patch_sdict, args.NUMPOINTS, VETOFNS, args.SAMPLER, useSEED)

    ## Write the points into the dir hierarchy
    ext = "" if len(patchDicts) == 1 else f".{npatch:02d}"
    for npt, p in enumerate(pps):
        npad = mkrunstr(npt)
        rundir = os.path.join(args.OUTDIR, npad+ext)

        ## Make the output dir
        if not os.path.exists(rundir):
            os.makedirs(rundir)

        ## Write out the params file
        if args.SAMPLER != "extend":
            writeParams(p, os.path.join(rundir, args.PFNAME))

        ## Instantiate template(s)
        p["N"] = npad #< Add the run number *after* writing out the params file
        writeTemplates(p, rundir, templates)

    ## Update the sample summary
    table += "#"
    for k in pps[0].keys():
        table += " %s" % k
    for point in pps:
        table += "\n"
        for i in point.values():
            try:
                table += "%e " % i
            except:
                table += "%s " % i

## Write the summary file
tpath = os.path.join(args.OUTDIR, "points.dat")
with open(tpath, "w") as pf:
    pf.write(table)
print(f"Sampled points summary written to {tpath}")
if args.DEBUG:
    print(table)
