import openpathsampling as paths
import numpy as np

from __future__ import print_function

#! lazy
import pyemma.coordinates as coor

#! lazy
ref_storage = paths.Storage('engine_store_test.nc', mode='r')

#! lazy
storage = paths.Storage('delete.nc', 'w')
storage.trajectories.save(ref_storage.trajectories[0])

"""
Explanation: PyEmma Featurizer Support
End of explanation
"""


def pyemma_generator(f):
    f.add_inverse_distances(f.pairs(f.select_Backbone()))

cv = paths.collectivevariable.PyEMMAFeaturizerCV(
    'pyemma', 
    pyemma_generator, 
    topology=ref_storage.snapshots[0].topology
).with_diskcache()

"""
Explanation: Import a PyEmma Coordinates Module
Using of pyemma featurizers or general other complex code requires a little trick to be storable. Since storing of code only works if we are not dependend on the context (scope) we need to wrap the construction of our featurizer in a function, that gets all it needs from the global scope as a parameter
End of explanation
"""


cv(ref_storage.trajectories[0]);

"""
Explanation: Now use this featurizer generating function to build a collective variable out of it. All we need for that is a name as usual, the generating function, the list of parameters - here only the topology and at best a test snapshot, a template.
End of explanation
"""


#! lazy
print(storage.save(cv))

"""
Explanation: Let's save it to the storage
End of explanation
"""


cv(storage.trajectories[0]);

"""
Explanation: and apply the featurizer to a trajectory
End of explanation
"""


cv(storage.snapshots.all());

py_cv = storage.cvs['pyemma']

store = storage.stores['cv%d' % storage.idx(py_cv)]
nc_var = store.variables['value']

assert(nc_var.shape[1] == 15)
print(nc_var.shape[1])

assert(nc_var.var_type == 'numpy.float32')
print(nc_var.var_type)

#! ignore
print(storage.variables['cvs_json'][:])

py_cv_idx = storage.idx(py_cv)
print(py_cv_idx)
py_emma_feat = storage.vars['attributes_json'][py_cv_idx]

erg = py_emma_feat(storage.snapshots);

#! lazy
print(erg[:,2:4])

storage.close()
ref_storage.close()

#! lazy
storage = paths.Storage('delete.nc', 'r')

cv = storage.cvs[0]

"""
Explanation: Sync to make sure the cache is written to the netCDF file.
End of explanation
"""


assert np.allclose(erg, cv(storage.snapshots))

storage.close()

"""
Explanation: Make sure that we get the same result
End of explanation
"""


