#!/usr/bin/env python
"""
This script takes some uq parameters from the command line
and starts up a PUQ session using them.  The RapptureHost()
class will not actually execute any jobs.  Job parameters
are instead saved to a CSV file named params[pid].csv.
PUQ status is saved into an hdf5 file, as usual.
"""

import sys
import os
from puq import NormalParameter, UniformParameter, Sweep, RapptureHost, Smolyak
import numpy as np
from puq.jpickle import unpickle

# Redirect stdout and stderr to files for debugging.
sys.stdout = open("uq_debug.out", 'w')
sys.stderr = open("uq_debug.err", 'w')

print sys.argv
pid, varlist, uq_type, args = sys.argv[1:]

dfile = "driver%s.xml" % pid
cvsname = "params%s.csv" % pid
hname = "puq_%s.hdf5" % pid

varlist = unpickle(varlist)
print "varlist=", varlist

v = {}
for p in varlist:
    name, dist = p
    name = str(name)
    print "name=", name
    print "dist=", dist
    if dist[0] == u'gaussian':
        v[name] = NormalParameter(name, name, mean=dist[1], dev=dist[2])
    elif dist[0] == u'uniform':
        v[name] = UniformParameter(name, name, min=dist[1], max=dist[2])
    else:
        print "ERROR: Unknown distribution type: %s" % dist[0]
        sys.exit(1)
if uq_type == "smolyak":
    uq = Smolyak(v.values(), args)
else:
    print "ERROR: Unknown UQ type: %s" % uq_type
    os.chdir('..')
    sys.exit(1)

# save parameter values to CSV file
vals = np.column_stack([p.values for p in uq.params])
names = ['@@'+str(p.name) for p in uq.params]
np.savetxt(cvsname, vals, delimiter=',', header=','.join(names), comments='')

# This just saves PUQ state into HDF5 file, so later we can have PUQ analyze
# the results and it knows about the input parameters, UQ method, etc.
sw = Sweep(uq, RapptureHost('puq', dfile), None)
sw.run(hname, overwrite=True)

print "Finished with get_params.py\n"