%matplotlib inline
import openpathsampling as paths

"""
Explanation: This is file runs the main calculation for the flexible length TPS simulation. It requires the file alanine_dipeptide_tps_equil.nc, which is written in the notebook alanine_dipeptide_tps_first_traj.ipynb.
In this file, you will learn:
* how to set up and run a flexible length TPS simulation
NB: This is a long calculation. In practice, it would be best to export the Python from this notebook, remove the live_visualizer, and run non-interactively on a computing node.
End of explanation
"""


old_storage = paths.Storage("tps_nc_files/alanine_dipeptide_tps_equil.nc", "r")

engine = old_storage.engines[0]
C_7eq = old_storage.volumes.find('C_7eq')
alpha_R = old_storage.volumes.find('alpha_R')
traj = old_storage.samplesets[len(old_storage.samplesets)-1][0].trajectory
phi = old_storage.cvs.find('phi')
psi = old_storage.cvs.find('psi')
template = old_storage.snapshots[0]

print engine.name

"""
Explanation: Load engine, trajectory, and states from file
End of explanation
"""


network = paths.TPSNetwork(C_7eq, alpha_R)

scheme = paths.OneWayShootingMoveScheme(network, selector=paths.UniformSelector(), engine=engine)

initial_conditions = scheme.initial_conditions_from_trajectories(traj)

storage = paths.Storage("tps_nc_files/alanine_dipeptide_tps.nc", "w", template)
sampler = paths.PathSampling(storage=storage,
                             move_scheme=scheme,
                             sample_set=initial_conditions).named("Flexible_TPS_Sampling")

"""
Explanation: TPS
As always, the process for setting up a simulation is:

Create a network
Create a move_scheme
Set up initial_conditions
Create the PathSampling object and run it.

Since we already created all the input to these when we set up the first trajectory, we can load use the versions we loaded above.
End of explanation
"""


#sampler.live_visualizer = paths.StepVisualization2D(network, phi, psi, [-3.14, 3.14], [-3.14, 3.14])
sampler.run(10000)

"""
Explanation: Note: 10000 steps will take a long time. If you just want to run a little bit, reduce this number.
End of explanation
"""


