%matplotlib inline
import matplotlib.pyplot as plt
import openpathsampling as paths
import numpy as np

%%time
storage = paths.AnalysisStorage('mstis_data.nc')

"""
Explanation: Analyzing a split MSTIS simulation
Included in this notebook:

Opening split files and look at the data
End of explanation
"""


mstis = storage.networks.load(0)

mstis.hist_args['max_lambda'] = { 'bin_width' : 0.02, 'bin_range' : (0.0, 0.5) }
mstis.hist_args['pathlength'] = { 'bin_width' : 5, 'bin_range' : (0, 150) }

%%time
mstis.rate_matrix(storage.steps, force=True)

"""
Explanation: Analyze the rate with no snapshots present in the analyzed file
End of explanation
"""


scheme = storage.schemes[0]

scheme.move_summary(storage.steps)

"""
Explanation: Move scheme analysis
End of explanation
"""


import openpathsampling.visualize as vis
reload(vis)
from IPython.display import SVG

tree = vis.PathTree(
    storage.steps[0:200],
    vis.ReplicaEvolution(replica=2, accepted=False)
)

SVG(tree.svg())

decorrelated = tree.generator.decorrelated
print "We have " + str(len(decorrelated)) + " decorrelated trajectories."

"""
Explanation: Replica move history tree
End of explanation
"""


from toy_plot_helpers import ToyPlot
background = ToyPlot()
background.contour_range = np.arange(-1.5, 1.0, 0.1)
background.add_pes(storage.engines[0].pes)

xval = paths.FunctionCV("xval", lambda snap : snap.xyz[0][0])
yval = paths.FunctionCV("yval", lambda snap : snap.xyz[0][1])
live_vis = paths.StepVisualizer2D(mstis, xval, yval, [-1.0, 1.0], [-1.0, 1.0])
live_vis.background = background.plot()

"""
Explanation: Visualizing trajectories
End of explanation
"""


storage.cvs

fallback = paths.Storage('mstis_traj.nc', 'r')

storage.fallback = fallback

live_vis.draw_samples(list(tree.samples))

"""
Explanation: to make this work we need the actual snapshot coordinates! These are not
present in the data file anymore so we attach the traj as a fallback.
We are not using analysis storage since we do not cache anything.
End of explanation
"""


