#!/usr/bin/env python
from argparse import ArgumentParser
import os, sys
import numpy as np
from matplotlib import pyplot as pp
from msmbuilder.Serializer import Serializer
from euclid.drift import drift
sys.path.append(os.path.join(os.path.abspath(os.path.dirname(__file__)), '../lib'))
from metrics import DWMetric
from Propagator import MixedHeight2DPropagator


def main():
    parser = ArgumentParser(description="Run embedded double well calcs")
    parser.add_argument('barrier1', type=float)
    parser.add_argument('barrier2', type=float)
    parser.add_argument('tau_max', type=int)
    parser.add_argument('tau_step', type=int, default=1)
    parser.add_argument('traj_length', type=int)
    parser.add_argument('output_dir', type=str, help='directory to save output')
    
    args = parser.parse_args()
    taus = [i for i in range(1, args.tau_max, args.tau_step)]
    if not os.path.exists(args.output_dir):
        os.mkdir(args.output_dir)
    
    def output_path(i):
        return os.path.join(os.path.abspath(args.output_dir), '%d.h5' % i)
    
    prop = MixedHeight2DPropagator(args.barrier1, args.barrier2)
    prop.run(args.traj_length)
    
    # save trajectory to disk
    Serializer(prop.trajectory).SaveToHDF('%s/trajectory.hdf' % args.output_dir)
    
    for i in range(2):
        drifts = drift(prop.trajectory, taus, DWMetric(i))
        entry = {'Data': drifts,
                 'Taus': taus,
                 'type': 'Dimension %s of %s' % (i, 2)
                }
        s = Serializer(entry).SaveToHDF(output_path(i))

if __name__ == '__main__':
    main()