#!/usr/bin/env python

import optparse
op = optparse.OptionParser(usage="%prog <datadir>")
opts, args = op.parse_args()

## The mandatory argument is the name of the data dir containing params-*.dat and out-*.yoda files
assert args
DATADIR = args[0]

## Read each params file, and the YODA file corresponding to it, to fill the param point and value dicts
# TODO: Provide Python lib helpers for reading param dirs in either flat or hier format
PARAMPOINTS, VALS = {}, {}
import os, glob
for pfile in glob.glob(os.path.join(DATADIR, "params-*.dat")):
    runstr = os.path.splitext(os.path.basename(pfile))[0].replace("params-", "")
    run = int(runstr)

    ## Get params for this run
    params = {}
    with open(pfile, "r") as pf:
        for l in pf:
            pname, pval = l.strip().split()
            params[pname] = float(pval)
    # print params.keys()

    ## Get data for this param point
    import yoda
    yfile = os.path.join(DATADIR, "out-%s.yoda" % runstr)
    if not os.path.exists(yfile):
        continue #< Skip this loop iteration if the YODA file doesn't exist (e.g. failed run)
    aos = yoda.read(yfile)
    val = aos["/ATLAS_2010_S8918562/d12-x01-y01"].bins[4].height
    # TODO: make this more complex, e.g. a weighted chi2 across lots of bins
    del aos # just to make sure this memory is really cleaned up asap

    ## Everything worked, so store this run params and value
    VALS[run] = val
    PARAMPOINTS[run] = params

## Convert the dicts to tuple views
PNAMES = sorted(PARAMPOINTS.keys())
PARAMS_VALS = [(PARAMPOINTS[pn].values(), VALS[pn]) for pn in PNAMES]
P, V = zip(*PARAMS_VALS)
# from pprint import pprint
# pprint(PARAMS_VALS)
# print P; print V

## Construct the interpolation
from professor2 import Ipol
ipol = Ipol(P, V, 5)

## Build plotting array of ipol values
import numpy as np
X = np.linspace(1, 5, 25)
Y = np.linspace(1, 3, 25)
M_X, M_Y = np.meshgrid(X, Y)
M_Z = np.empty(M_X.shape)
for i in xrange(M_X.shape[0]):
    for j in xrange(M_X.shape[1]):
        M_Z[i][j] = ipol.value(M_X[i][j], M_Y[i][j])

## Plot as a 3D view of the 2D data
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
X1, Y1 = zip(*P)
Z1 = V
ax.scatter(X1, Y1, Z1, c="r")
ax.plot_wireframe(M_X, M_Y, M_Z, rstride=1, cstride=1)
ax.set_xlabel("PT0REF")
ax.set_ylabel("EXPPOW")
plt.show()
