#!/usr/bin/env python
# -*- coding: utf-8 -*-

if __name__ == '__main__':

    import argparse

    parser = argparse.ArgumentParser(
        description="Python wrapper example used for the PRACE training on HPC and uncertainty.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument('-tmp', default='/tmp', type=str,
        help='Root directory on which temporary working directories will be' +
             'created for each independent simulation.')

    parser.add_argument('-seed', default=int(0), type=int,
        help='Seed number for the random number generator')

    parser.add_argument('-MonteCarlo', nargs=1,
        help="Launch a MonteCarlo simulation of given size")

    parser.add_argument('-X', nargs='*',
        help='List of floats [X1, X2.. Xp] or PATH to a pickled DOE')

    parser.add_argument('-n_cpus', default=-1, type=int,
        help="(Optional) number of cpus to use.")

    parser.add_argument('-backend', default='joblib', type=str,
        choices=['joblib', 'multiprocessing', 'ipyparallel'],
        help="Whether to parallelize using 'joblib' or 'multiprocessing'.")

    parser.add_argument('-run', default=False, type=bool, nargs='?',
        const='True', help='If True, run the model', choices=[True, False])

    parser.add_argument('-dump', default=False, type=bool, nargs='?',
        const='True', choices=[True, False],
        help='If True, dump the output for later posttreatment')

    args = parser.parse_args()

    import otwrapy as otw
    import openturns as ot
    from otwrapy.examples.beam import X_distribution, Wrapper

    model = otw.Parallelizer(Wrapper(tmpdir=args.tmp, sleep=1),
        backend=args.backend, n_cpus=args.n_cpus)

    print("The wrapper has been instantiated as 'model'.")

    if args.MonteCarlo is not None:
        ot.RandomGenerator.SetSeed(args.seed)
        N = int(args.MonteCarlo[0])
        X = X_distribution.getSample(N)
        print("Generated a MonteCarlo DOE of size {}".format(N))

    elif args.X is not None:
        if isinstance(args.X[0], str) and os.path.isfile(args.X[0]):
            X = otw.load_array(args.X[0])
            print("Loaded a DOE of size {} from file: '{}'".format(X.getSize(),
                args.X[0]))
        else:
            X = ot.Point([float(x) for x in args.X])

    if args.run:
        Y = model(X)
        # Dump the results if asked
        if args.dump:
            otw.dump_array(Y, 'OutputSample.pkl')
            print("The output has been saved to 'OutputSample.pkl'")
        else:
            print("Finished evaluationg the model. Take a look at 'Y' variable.")
    elif (args.MonteCarlo is not None) or (args.X is not None):
        print("The desired input is ready to be run using --> 'Y = model(X)'")
