#!/usr/bin/env python
import os, sys
import numpy as np
from argparse import ArgumentParser
from msmbuilder.Serializer import Serializer
sys.path.append(os.path.join(os.path.abspath(os.path.dirname(__file__)), '../lib'))
from metrics import DWMetric
import matplotlib.pyplot as pp
from PMF import plotpmf_gkde, plotpmf_hist
from scipy.stats import gaussian_kde
    
    
def main():
    parser = ArgumentParser(description='Plot PMF')
    parser.add_argument('traj_file', type=str)
    parser.add_argument('metrics', type=int, nargs='+', help='only DWMetrics at this time')
    parser.add_argument('-nbins', type=int)
    parser.add_argument('-gkde', action='store_true', default=False)
    parser.add_argument('-hist', action='store_true', default=False)
    args = parser.parse_args()
    if args.nbins is not None:
        args.hist = True
    
    traj = Serializer.LoadFromHDF(args.traj_file)
    metrics = [DWMetric(i) for i in args.metrics]
    
    datas = [traj['XYZList'][:, i] for i in range(len(metrics))]
    
    
    if args.gkde:
        pp.figure()
        plotpmf_gkde(datas)
    if args.hist:
        pp.figure()
        plotpmf_hist(datas, args.nbins)
    pp.show()
        
    
if __name__ == '__main__':
    main()