%matplotlib inline

import openpathsampling as paths
import numpy as np
import math

# the openpathsampling OpenMM engine
import openpathsampling.engines.openmm as eng

"""
Explanation: Run from bootstrap paths
Now we will use the initial trajectories we obtained from bootstrapping to run an MSTIS simulation. This will show both how objects can be regenerated from storage and how regenerated equivalent objects can be used in place of objects that weren't stored.
Tasks covered in this notebook:
* Loading OPS objects from storage
* Ways of assigning initial trajectories to initial samples
* Setting up a path sampling simulation with various move schemes
* Visualizing trajectories while the path sampling is running
End of explanation
"""


old_store = paths.AnalysisStorage("ala_mstis_bootstrap.nc")

"""
Explanation: Loading things from storage
First we'll reload some of the stuff we stored before. Of course, this starts with opening the file.
End of explanation
"""


print "PathMovers:", len(old_store.pathmovers)
print "Engines:", len(old_store.engines)
print "Samples:", len(old_store.samples)
print "Trajectories:", len(old_store.trajectories)
print "Ensembles:", len(old_store.ensembles)
print "SampleSets:", len(old_store.samplesets)
print "Snapshots:", len(old_store.snapshots)
print "Networks:", len(old_store.networks)

"""
Explanation: A lot of information can be recovered from the old storage, and so we don't have the recreate it. However, we did not save our network, so we'll have to create a new one. Since the network creates the ensembles, that means we will have to translate the trajectories from the old ensembles to new ensembles.
End of explanation
"""


# template = old_store.snapshots[0]
engine = old_store.engines['default']
mstis = old_store.networks[0]
sset = old_store.tag['sampleset']

"""
Explanation: Loading from storage is very easy. Each store is a list. We take the 0th snapshot as a template (it doesn't actually matter which one) for the next storage we'll create. There's only one engine stored, so we take the only one.
End of explanation
"""


platform = 'CUDA'
engine.initialize(platform)

print 'Engine uses platform `%s`' % engine.platform

sset.sanity_check()

"""
Explanation: initialize engine
if we do not select a platform the fastest possible will be chosen but we explicitly request to use the one in the config file
End of explanation
"""


# logging creates ops_output.log file with details of what the calculation is doing
#import logging.config
#logging.config.fileConfig("logging.conf", disable_existing_loggers=False)

storage = paths.storage.Storage("ala_mstis_production.nc", "w")

storage.snapshots.save(old_store.snapshots[0]);

"""
Explanation: Running RETIS
Now we run the full calculation. Up to here, we haven't been storing any of our results. This time, we'll start a storage object, and we'll save the network we've created. Then we'll run a new PathSampling calculation object.
End of explanation
"""


scheme = paths.DefaultScheme(mstis, engine)

"""
Explanation: Before we can sample we still need to set the actual MoveScheme which determines the
set of moves to apply to our set of samples and effectively doing the steps in
replica (sampleset) space. We pick the default scheme for mstis and feed it with
the engine to be used.
End of explanation
"""


mstis_calc = paths.PathSampling(
    storage=storage,
    sample_set=sset,
    move_scheme=scheme
)
mstis_calc.save_frequency = 10

"""
Explanation: and finally generate the PathSampler object to conduct the simulation.
End of explanation
"""


mstis_calc.run(5)

print len(storage.steps)

# commented out during development, so we can "run all" and then do more
storage.close()

"""
Explanation: Now everything is ready: let's run the simulation! The first step takes a little since all
necessary information, i.e. the engines, topologies, initial snapshots, ..., need to be
stored. Then the monte carlo steps will be performed.
End of explanation
"""


