import openpathsampling as paths
import openmmtools as omt
import simtk.openmm as omm
import simtk.unit as u
import mdtraj as md

import openpathsampling.engines.openmm as eng

from __future__ import print_function

"""
Explanation: OpenMMTools Integration
This enables OPS to use any testsystem from openmmtools to be used as well as add ANY possible OpenMM setup, not just the ones that we provided before. The old OpenMMEngine has been rename intp SimpleOpenMMEngine
End of explanation
"""


testsystem = omt.testsystems.AlanineDipeptideVacuum()

"""
Explanation: Create Alanine in Vacuum and run it using OPS.
End of explanation
"""


#! skip
{ key: type(value) for key, value in testsystem.__dict__.items()}

"""
Explanation: Let's have a look at the content
End of explanation
"""


template = eng.snapshot_from_testsystem(testsystem)

"""
Explanation: An OpenMM simulation in OPS needs 3 ingredients to function
1. A template snapshot (which contains topolgy and coordinates/velocities), 
2. an OpenMM system object, and 
3. an OpenMM integrator object.
In the following we will get these parts.
1. The template
The information for the template is contained in the _topology and the _positions object.
End of explanation
"""


system = testsystem.system

"""
Explanation: 2. The system
This is easy since we directly get it from the testsytem objects
End of explanation
"""


integrator = omm.VerletIntegrator(
    0.002 * u.picoseconds    
)

"""
Explanation: 3. The integrator
The openmmtools system is (almost) independent of the integrator and openmm provides us with lots of options. For now we pick something simple.
End of explanation
"""


engine = eng.Engine(
    template.topology, 
    system, 
    integrator)

"""
Explanation: build the engine
End of explanation
"""


traj = engine.generate(template, [paths.LengthEnsemble(10).can_append])

print(traj)

#! skip
print traj[5].coordinates

psi = md.compute_psi(traj.to_mdtraj())

#! ignore
psi[1][3:8]

st = paths.Storage('engine_store_test.nc', mode='w')

st.engines.save(engine);
st.tag['template'] = template

"""
Explanation: Let's run a simulation of 10 steps.
End of explanation
"""


#! skip
st.variables['engines_json'][0][0:256] + '...'

st.save(traj);

st.close()

"""
Explanation: The engine is stored by using the XML serialization of OpenMM and put this into the storage.
End of explanation
"""


st = paths.AnalysisStorage('engine_store_test.nc')

engine = st.engines[0]

"""
Explanation: try to load the stored engine
End of explanation
"""


template = st.tag['template']

traj = engine.generate(template, [paths.LengthEnsemble(10).can_append])

print(traj)

st.close()

"""
Explanation: Run another 10 steps to test
End of explanation
"""


st = paths.Storage('engine_store_test.nc', "a")
st.save(traj)
st.close()

"""
Explanation: And save the trajectory for future use (other tests)
End of explanation
"""


