filename = "alanine_dipeptide_tps.nc"

"""
Explanation: Analyzing the flexible path length simulation
Load the file, and from the file pull our the engine (which tells us what the timestep was) and the move scheme (which gives us a starting point for much of the analysis).
End of explanation
"""


%%time
flexible = paths.Storage(filename, mode='r')

engine = flexible.engines[0]
flex_scheme = flexible.schemes[0]

print("File size: {0} for {1} steps, {2} snapshots".format(
    flexible.file_size_str,
    len(flexible.steps),
    len(flexible.snapshots),
))

"""
Explanation: Opening a large storage file can take some time. In addition, the AnalysisStorage object tries to pre-cache things, which makes opening it take longer, but can make analysis faster.
End of explanation
"""


%%time
flex_scheme.move_summary(flexible.steps)

"""
Explanation: That tell us a little about the file we're dealing with. Note that the number of snapshots listed is twice the number of configurations. That's because every snapshot object stores a virtual version of its velocity-reversed copy. Both share the same configuration and velocity storage; one just flips the signs on the velocities.
Now we'll start analyzing the contents of that file. We used a very simple move scheme (only shooting), so the main information that the move_summary gives us is the acceptance of the only kind of move in that scheme. See the MSTIS examples for more complicated move schemes, where you want to make sure that frequency at which the move runs is close to what was expected.
End of explanation
"""


tree = ops_vis.PathTree(
    flexible.steps[0:21],
    ops_vis.ReplicaEvolution(
        replica=0
    )
)

with open("flexible_path_tree.svg", mode='w') as f:
    f.write(tree.svg())

SVG(tree.svg())

print("Decorrelated trajectories (first 20 steps):", len(tree.generator.decorrelated_trajectories))

%%time
full_tree = ops_vis.PathTree(
    flexible.steps,
    ops_vis.ReplicaEvolution(
        replica=0
    )
)
print("Total decorrelated trajectories:", len(full_tree.generator.decorrelated_trajectories))

"""
Explanation: Replica history tree and decorrelated trajectories
The PathTree object gives us both the history tree (often called the "move tree") and the number of decorrelated trajectories. It takes the history of a given replica (in our case, replica 0, the only one) and builds a description of the move history.
A PathTree is made for a certain set of Monte Carlo steps. First, we make a tree of only the first 20 steps in order to visualize it. (All 10000 steps would be unwieldy.) Note that only the accepted steps are shown by default (see examples for path tree visualization for customization details.)
After the visualization, we make a second PathTree of all the steps, in order to count the number of decorrelated trajectories.
End of explanation
"""


path_lengths = [len(step.active[0].trajectory) for step in flexible.steps]
plt.hist(path_lengths, bins=40, alpha=0.5);
plt.xlabel("Path length (frames)")
plt.ylabel("Number of paths")
#print("Maximum:", max(path_lengths), "("+str(max(path_lengths)*engine.snapshot_timestep)+")")
#print("Average:", "{0:.2f}".format(np.mean(path_lengths)), "("+(np.mean(path_lengths)*engine.snapshot_timestep).format("%.3f")+")")
plt.tight_layout()

"""
Explanation: Path length distribution
Flexible length TPS gives a distribution of path lengths. Here we calculate the length of every accepted trajectory, then histogram those lengths, and calculate the maximum and average path lengths.
We also use engine.snapshot_timestep to convert the count of frames to time, including correct units.
End of explanation
"""


from openpathsampling.numerics import HistogramPlotter2D

psi = flexible.cvs['psi']
phi = flexible.cvs['phi']
deg = 180.0 / np.pi

path_density = paths.PathDensityHistogram(cvs=[phi, psi],
                                          left_bin_edges=(-180/deg,-180/deg),
                                          bin_widths=(2.0/deg,2.0/deg))

%%time

path_dens_counter = path_density.histogram([s.active[0].trajectory for s in flexible.steps])

"""
Explanation: Path density histogram
Next we will create a path density histogram. Calculating the histogram itself is quite easy: first we reload the collective variables we want to plot it in (we choose the phi and psi angles). Then we create the empty path density histogram, by telling it which CVs to use and how to make the histogram (bin sizes, etc). Finally, we build the histogram by giving it the list of active trajectories to histogram.
End of explanation
"""


tick_labels = np.arange(-np.pi, np.pi+0.01, np.pi/4)
plotter = HistogramPlotter2D(path_density, 
                             xticklabels=tick_labels,
                             yticklabels=tick_labels, 
                             label_format="{:4.2f}")
ax = plotter.plot(cmap="Blues")

"""
Explanation: Now we've built the path density histogram, and we want to visualize it. We have a convenient plot_2d_histogram function that works in this case, and takes the histogram, desired plot tick labels and limits, and additional matplotlib named arguments to plt.pcolormesh.
End of explanation
"""


