#!/usr/bin/env python3

"""
Sample and instantiate templates.
"""

import numpy as np

def writeParams(P, templates, outdir, fname="params.dat"):
    from os.path import join, exists
    for num, p in enumerate(P):
        npad = "{}".format(num).zfill(1+int(np.ceil(np.log10(len(P)))))
        outd = join(outdir, npad)
        outf = join(outd, fname)

        if not exists(outd):
            import os
            os.makedirs(outd)

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

        ## Instantiate template(s)
        p["N"] = npad #< Add the run number *after* writing out the params file
        for tbasename, tmpl in templates.items():
            txt = tmpl.format(**p)
            tname = join(outd, tbasename)
            with open(tname, "w") as tf:
                tf.write(txt)

def sample(boxdef, npoints):
    is_json=True;
    import json
    with open(boxdef) as f:
        c = f.read(1)
        if c!="{": is_json=False

    if is_json:
        with open(boxdef) as f:
            import json
            B = json.load(f)
    else:
        with open(boxdef) as f:
            _=[l.strip().split() for l in f]
            B = {x[0]:[float(x[1]), float(x[2])] for x in _}

    # Alphabetical sort for ranges
    porder = sorted(B.keys())
    xmin = [B[x][0] for x in porder]
    xmax = [B[x][1] for x in porder]

    # Randomly sampled points
    # TODO add lhs as preferred option?
    RSP = np.random.uniform(low=xmin, high=xmax,size=(npoints, len(xmin)))
    for p in porder:
        print(p)

    # As dictionaries
    return [ dict(zip(porder, x)) for x in RSP]


if __name__=="__main__":
    import sys
    PP = sample(sys.argv[1], int(sys.argv[2]))

    # First arg: json with param box
    # 2nd arg: number of points to sample i.e. number of sub directories to create
    # 3rd arg: template

    import os
    tname = os.path.basename(sys.argv[3])
    TEMPLATES={}
    with open(sys.argv[3], "r") as f:
        TEMPLATES[tname] = f.read()

    writeParams(PP, TEMPLATES, "newscan")
