#!/usr/bin/env python
from argparse import ArgumentParser
import os, sys
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 EmbeddedDoubleWellPropagator as EDWProp


def main():
    parser = ArgumentParser(description="Run embedded double well calcs")
    parser.add_argument('num_dims', type=int)
    parser.add_argument('k', type=float)
    parser.add_argument('tau', 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, args.tau_step)]
    
    def output_path(i):
        if not os.path.exists(args.output_dir):
            os.mkdir(args.output_dir)
        return os.path.join(os.path.abspath(args.output_dir), '%d.h5' % i)
    
    prop = EDWProp(args.num_dims, args.k)
    prop.run(args.traj_length)
    for i in range(args.num_dims):
        drifts = drift(prop.trajectory, taus, DWMetric(i))
        entry = {'Data':drift(prop.trajectory, taus, DWMetric(i)),
                 'Taus': taus,
                 'type': 'Dimension %s of %s' % (i, args.num_dims)
                }
        s = Serializer(entry).SaveToHDF(output_path(i))
    
    
    
    
if __name__ == '__main__':
    main()