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


def main():
    parser = ArgumentParser(description="Evaluate hybrid")
    parser.add_argument('-d', '--directory', default='Observables',
                         help='Directory containing files with the drift under different metrics. Each file should end in the suffix ".h5"')
    parser.add_argument('weights', type=float, nargs='+')
    
    args = parser.parse_args()
    metric_fns = glob('%s/*.h5' % args.directory)
    num_metrics = len(metric_fns)
    if num_metrics == 0:
        parser.error('No metrics found')
    print 'Found %d metrics files: %s' % (num_metrics, str([os.path.split(fn)[1] for fn in metric_fns]))
    
    trajectory = Serializer.LoadFromHDF('%s/trajectory.hdf' % args.directory)

    if not len(args.weights) == num_metrics:
        parser.exit('Wrong number of weights. Should be %d' % num_metrics)
    base_metrics = [DWMetric(i) for i in range(num_metrics)]
    hybrid_metric = HybridMetric(base_metrics, args.weights)
    
    
    # plot autocorrelation
    for i in range(num_metrics):
        metric = DWMetric(i)
        ptraj = metric.prepare_trajectory(trajectory)
        ac = fft_acf(metric.one_to_all(ptraj, ptraj, 0))
        pp.semilogx(ac, label='metric %d' % i)
    ptraj = hybrid_metric.prepare_trajectory(trajectory)   
    ac = fft_acf(hybrid_metric.one_to_all(ptraj, ptraj, 0))
    pp.semilogx(ac, label='hybrid')
    pp.legend()
    pp.title('autocorrelation')
    pp.savefig('%s/ac.png' % args.directory)
    
    # plot drift time mean/std
    pp.figure()
    for i in range(num_metrics):
        m = Serializer.LoadFromHDF(metric_fns[i])
        drifts = m['Data']
        taus = m['Taus']
        ma = np.ma.masked_less(drifts, 0)
        ma /= np.mean(ma[-1,:])
        pp.subplot(211)
        pp.plot(np.sqrt(taus), np.mean(ma, axis=1), 'x-', label=metric_fns[i])
        pp.subplot(212)
        pp.plot(np.sqrt(taus), np.std(ma, axis=1), 'x-', label=metric_fns[i])
    hybrid_drifts = drift(trajectory, taus, hybrid_metric)
    ma = np.ma.masked_less(hybrid_drifts, 0)
    ma /= np.mean(ma[-1, :])
    pp.subplot(211); pp.title('mean drift progress')
    pp.plot(np.sqrt(taus), np.mean(ma, axis=1), 'x-', label='hybrid')
    pp.legend(loc=4)
    pp.xlabel('sqrt t')
    pp.subplot(212); pp.title('std drift progress')
    pp.plot(np.sqrt(taus), np.std(ma, axis=1), 'x-', label='hybrid')
    pp.xlabel('sqrt t')
    pp.legend(loc=4)
    pp.savefig('%s/drift.png' % args.directory)
    
    #pp.show()

if __name__ == '__main__':
    main()