#!/usr/bin/env python
"""
Once submit has finished with the jobs, this function is called to have PUQ
process the results.
"""
from __future__ import print_function
import sys
import os
import numpy as np
import h5py
import re
import puq
from puq.jpickle import unpickle
import xml.etree.ElementTree as xml
import Rappture
import StringIO

# Redirect stdout and stderr to files for debugging.
# Append to the files created in get_params.py
sys.stdout = open("uq_debug.out", 'a')
sys.stderr = open("uq_debug.err", 'a')


# Restore the state of a PUQ session from a HDF5 file.
def load_from_hdf5(name):
    h5 = h5py.File(name, 'r+')
    sw = unpickle(h5['private/sweep'].value)
    sw.fname = os.path.splitext(name)[0]
    h5.close()

    sw.psweep._sweep = sw

    if hasattr(sw.psweep, 'reinit'):
        sw.psweep.reinit()
    return sw


# Plots probability curves
def plot_pdf_curve(xvals, vname, percent, desc):
    print('plot_pdf_curve %s %s %s' % (vname, percent, desc))
    # compute upper and lower percentiles
    pm = (100 - percent)/200.0
    pp = 1 - pm

    # collect data into an array
    xarr = np.empty(len(xvals[vname]))
    yp = np.empty(len(xvals[vname]))
    ym = np.empty(len(xvals[vname]))
    for vindex in sorted(xvals[vname].keys()):
        xarr[vindex] = xvals[vname][vindex]
        yp[vindex] = pcurves[vname][vindex].ppf(pp)
        ym[vindex] = pcurves[vname][vindex].ppf(pm)

    curve = xml.SubElement(dout, 'curve', {'id': 'curve_pdf-%s-%s' % (vname, percent)})
    about = xml.SubElement(curve, 'about')
    if percent == 0:
        xml.SubElement(about, 'label').text = "mean"
    else:
        xml.SubElement(about, 'label').text = "middle %s%%" % percent
    xml.SubElement(about, 'group').text = '%s (Probability)' % desc
    comp = xml.SubElement(curve, 'component')
    xy = xml.SubElement(comp, 'xy')
    pts = ""
    for x, y in zip(xarr, yp):
        pts += "%s %s " % (x, y)
    if percent == 0:
        pts += '\n'
    else:
        for x, y in reversed(zip(xarr, ym)):
            pts += "%s %s " % (x, y)
        pts += "%s %s\n" % (xarr[0], yp[0])
    xy.text = pts


def dist(x1, y1, x2, y2, x3, y3): # x3,y3 is the point
    diffx = x2-x1
    diffy = y2-y1
    sqdiff = float(diffx**2 + diffy**2)

    u = ((x3 - x1) * diffx + (y3 - y1) * diffy) / sqdiff

    if u > 1:
        u = 1
    elif u < 0:
        u = 0

    x = x1 + u * diffx
    y = y1 + u * diffy

    dx = x - x3
    dy = y - y3

    return np.sqrt(dx*dx + dy*dy)


# Plots advanced probability curves
def plot_pdf_acurve(hf, acurves, vname, percent, desc):
    print('plot_pdf_acurve %s %s %s' % (vname, percent, desc))
    # compute upper and lower percentiles
    pm = (100 - percent)/200.0
    pp = 1 - pm
    vlen = len(acurves[vname])

    print('vlen=', vlen)
    # collect data into an array
    xp = np.empty(vlen)
    xm = np.empty(vlen)
    yp = np.empty(vlen)
    ym = np.empty(vlen)

    for vindex in sorted(acurves[vname].keys()):
        print('vindex=', vindex)
        x1 = acurves[vname][vindex][0].ppf(pp)
        x2 = acurves[vname][vindex][0].ppf(pm)
        y1 = acurves[vname][vindex][1].ppf(pp)
        y2 = acurves[vname][vindex][1].ppf(pm)
        xd = hf['/output/data/%s.x[%d]' % (vname, int(vindex))].value
        yd = hf['/output/data/%s.y[%d]' % (vname, int(vindex))].value
        #p1, p2 = get_closest2()
        if int(vindex) == 2:
            print (x1, x2, y1, y2)
            print ('xd=', repr(xd))
            print ('yd=', repr(yd))

        #xp[vindex] = x1
        #xm[vindex] = x2
        #yp[vindex] = y1
        #ym[vindex] = y2
    return
    curve = xml.SubElement(dout, 'curve', {'id': 'curve_pdf-%s-%s' % (vname, percent)})
    about = xml.SubElement(curve, 'about')
    if percent == 0:
        xml.SubElement(about, 'label').text = "mean"
    else:
        xml.SubElement(about, 'label').text = "middle %s%%" % percent
    xml.SubElement(about, 'group').text = '%s (Probability)' % desc
    comp = xml.SubElement(curve, 'component')
    xy = xml.SubElement(comp, 'xy')
    pts = ""
    for x, y in zip(xarr, yp):
        pts += "%s %s " % (x, y)
    if percent == 0:
        pts += '\n'
    else:
        for x, y in reversed(zip(xarr, ym)):
            pts += "%s %s " % (x, y)
        pts += "%s %s\n" % (xarr[0], yp[0])
    xy.text = pts


def plot_pdf(v, pdf, desc):
    id = '%s (PDF)' % desc
    elem = xml.SubElement(dout, 'curve', {'id': id})
    about = xml.SubElement(elem, 'about')
    xml.SubElement(about, 'label').text = id
    yaxis = xml.SubElement(elem, 'yaxis')
    xml.SubElement(yaxis, 'label').text = 'Probability'

    component = xml.SubElement(elem, 'component')
    xy = xml.SubElement(component, 'xy')
    pts = "%s 0\n" % pdf.x[0]
    for x, y in zip(pdf.x, pdf.y):
        pts += "%s %s\n" % (x, y)
    pts += "%s 0\n" % pdf.x[-1]
    xy.text = pts

# Plotting response distance
# plot_resp1 distance
# Plotting response maxheight
# plot_resp1 maxheight
# desc= The distance the projectile travels.
# label= Distance

# v=-0.610951707827365*angle**2 + 54.985653704463*angle - 217.610577363519

# SURROGATE MODEL ERROR: 0.442%
# desc= The maximum height the projectile reaches.
# label= Maximum Height

# v=3.42068544561648e-12*angle**2 + 8.66308066792632*angle - 134.909576819562

# SURROGATE MODEL ERROR: 0.586%
# [[  45.          254.92905324]
#  [  30.19262971  128.95166845]
#  [  59.80737029  380.90643804]
#  [  34.52960806  163.81660645]
#  [  55.47039194  346.04150004]]


def write_responses(io, h5):
    uqtype = h5.attrs['UQtype']
    for v in h5[uqtype]:
        if '[' in v:
            # It is a curve. Ignore.
            continue
        try:
            desc = h5['%s/%s' % (uqtype, v)].attrs['description']
        except:
            desc = ''
        try:
            label = h5['%s/%s' % (uqtype, v)].attrs['label']
        except:
            label = ''

        rsp = h5['/%s/%s/response' % (uqtype, v)].value

        rout = io['output.response(%s)' % label]
        rout['value'] = rsp
        rout['about.description'] = desc
        rout['about.label'] = label + " (Response)"

        rs = unpickle(rsp)
        rout['variables'] = ' '.join([str(p.name) for p in rs.params])
        labels = ' '.join([repr(str(p.label)) for p in rs.params])
        rout['labels'] = labels.replace("'", '"')

        if type(rs) == puq.response.ResponseFunc:
            rout['equation'] = rs.eqn
            rout['rmse'] = "{:6.3g}".format(rs.rmse()[1])

        rout['data'] = rs.data


def write_params(h5, out):
    params = map(str, h5['/input/params'].keys())
    print('#' * 80, file=out)
    print('INPUT PARAMETERS', file=out)

    for pname in params:
        print('-' * 80, file=out)
        p = puq.unpickle(h5['/input/params/' + pname].value)
        cname = p.__class__.__name__[:-9]
        pdf_str = '%s [%s - %s] mean=%s dev=%s mode=%s' % (cname, p.pdf.range[0], p.pdf.range[1], p.pdf.mean, p.pdf.dev, p.pdf.mode)

        print("Name:", p.name, file=out)
        try:
            print("Label:", p.label, file=out)
        except:
            pass
        print("Desc:", p.description, file=out)
        print('Value:', pdf_str, file=out)
    print('#' * 80, file=out)
    print(file=out)


def write_summary(io, h5):
    outstr = StringIO.StringIO()
    write_params(h5, outstr)
    uqtype = h5.attrs['UQtype']
    for v in h5[uqtype]:
        if '[' in v:
            # It is a curve. Ignore.
            continue
        desc = h5['%s/%s' % (uqtype, v)].attrs['description']
        print("QoI: %s (%s)" % (v, desc), file=outstr)
        rs = unpickle(h5['/%s/%s/response' % (uqtype, v)].value)
        if type(rs) == puq.response.ResponseFunc:
            print("\nv=%s\n" % rs.eqn, file=outstr)
            print("SURROGATE MODEL ERROR:{:6.3g}%".format(rs.rmse()[1]), file=outstr)
        sens = puq.unpickle(h5['/%s/%s/sensitivity' % (uqtype, v)].value)
        max_name_len = max(map(len, [p[0] for p in sens]))
        print("\nSENSITIVITY:", file=outstr)
        print("Var%s     u*          dev" % (' '*(max_name_len)), file=outstr)
        print('-'*(28+max_name_len), file=outstr)
        for item in sens:
            pad = ' '*(max_name_len - len(item[0]))
            print("{}{}  {:10.4g}  {:10.4g}".format(pad, item[0],
                item[1]['ustar'], item[1]['std']), file=outstr)
        print('-'*(28+max_name_len), file=outstr)
        print(file=outstr)
    iostr = io['output.string(UQ Summary)']
    iostr['about.label'] = 'UQ Summary'
    iostr['current'] = outstr.getvalue()
    outstr.close()

sw = load_from_hdf5(sys.argv[1])
sw.analyze()

h5 = h5py.File(sys.argv[1], 'r+')
dtree = xml.parse('run_uq.xml')
droot = dtree.getroot()
dout = droot.find('output')

# curves built from pdfs
pcurves = {}
xvals = {}
acurves = {}

reg1 = re.compile('([ \da-zA-Z]+)\[([ \d]+)\]')
reg2 = re.compile('([ \da-zA-Z]+)\.([xy])\[([ \d]+)\]')

uqtype = h5.attrs['UQtype']
for v in h5[uqtype]:
    rsp = h5['/%s/%s/response' % (uqtype, v)].value
    print("rsp=", rsp)
    rs = unpickle(rsp)
    pdf = rs.pdf(fit=False)
    odata = h5['/output/data/%s' % v]

    # For curves built from pdfs, just put them in a dict for now
    if 'x' in odata.attrs:
        matches = reg1.findall(v)
        vname, vindex = matches[0]
        vindex = int(vindex)
        if vname not in pcurves:
            pcurves[vname] = {}
            xvals[vname] = {}
        xvals[vname][vindex] = odata.attrs['x']
        pcurves[vname][vindex] = pdf
    elif reg2.findall(v) != []:
        matches = reg2.findall(v)
        vname, xy, vindex = matches[0]
        if vname not in acurves:
            acurves[vname] = {}
        if vindex not in acurves[vname]:
            acurves[vname][vindex] = {}
        if xy == 'x':
            acurves[vname][vindex][0] = pdf
        else:
            acurves[vname][vindex][1] = pdf
    else:
        desc = h5['/output/data/%s' % v].attrs['label']
        plot_pdf(v, pdf, desc)

# now do pdf curves
for vname in xvals:
    desc = h5['/output/data/%s[0]' % vname].attrs['description']
    plot_pdf_curve(xvals, vname, 0, desc)
    plot_pdf_curve(xvals, vname, 50, desc)
    plot_pdf_curve(xvals, vname, 95, desc)

# for vname in acurves:
#    desc = h5['/output/data/%s.x[0]' % vname].attrs['description']
#    plot_pdf_acurve(h5, acurves, vname, 0, desc)
#    plot_pdf_acurve(h5, acurves, vname, 50, desc)
#    plot_pdf_acurve(h5, acurves, vname, 95, desc)


# If more than one variable, display sensitivity.
# Curves have indexed variables, so skip them.
if len(h5['/input/params']) > 1 and ['[' in x for x in h5[uqtype]].count(False):
    for v in h5[uqtype]:
        if '[' in v:
            continue
        desc = h5['/output/data/%s' % v].attrs['label']
        sens = unpickle(h5['/%s/%s/sensitivity' % (uqtype, v)].value)
        elem = xml.SubElement(dout, 'histogram', {'id': 'sens-%s' % v})
        about = xml.SubElement(elem, 'about')
        xml.SubElement(about, 'label').text = '%s (Sensitivity)' % desc
        xml.SubElement(elem, 'type').text = 'scatter'
        xaxis = xml.SubElement(elem, 'xaxis')
        xml.SubElement(xaxis, 'label').text = 'Parameters'
        yaxis = xml.SubElement(elem, 'yaxis')
        xml.SubElement(yaxis, 'label').text = 'Sensitivity'
        comp = xml.SubElement(elem, 'component')
        xy = xml.SubElement(comp, 'xy')
        pts = ''
        for name in sens:
            n = name[0]
            try:
                n = h5['/input/params/%s' % n].attrs['label']
            except:
                pass
            pts += "\"%s\" %s\n" % (n, name[1]['ustar'])
        xy.text = pts

with open('run_uq.xml', 'w') as f:
    f.write("<?xml version=\"1.0\"?>\n")
    dtree.write(f)

io = Rappture.PyXml('run_uq.xml')
write_responses(io, h5)
write_summary(io, h5)
io.close()

h5.close()