import molsysmt as msm
from molsysmt import pyunitwizard as puw
import numpy as np

import openmm as mm
from openmm import app
from openmm import unit

#import os
#from sys import stdout
#from mdtraj.reporters import HDF5Reporter
#from pathlib import Path
#import shutil

modeller = msm.convert(msm.systems['Trp-Cage']['1l2y.h5msm'], to_form='openmm.Modeller', structure_indices=0)
forcefield = app.ForceField("amber14-all.xml", "amber14/tip3p.xml")
system = forcefield.createSystem(modeller.topology, nonbondedMethod=app.NoCutoff, constraints=app.HBonds)
integrator = mm.LangevinIntegrator(300*unit.kelvin, 1.0/unit.picosecond, 2.0*unit.femtoseconds)
platform = mm.Platform.getPlatformByName('CPU')
simulation = app.Simulation(modeller.topology, system, integrator, platform)
simulation.context.setPositions(modeller.positions)
simulation.minimizeEnergy()
simulation.context.setVelocitiesToTemperature(300*unit.kelvin)

reporter = msm.thirds.openmm.reporters.StructuresDictReporter(100, time=True, coordinates=True, velocities=True,
             potentialEnergy=True, kineticEnergy=True, totalEnergy=True, temperature=True, box=True)

simulation.reporters.append(reporter)

simulation.step(1000)

traj_dict=reporter.finalize()

traj_dict

{'time': [0.  0.2 0.4 0.6 0.8 1.  1.2 1.4 1.6 1.8 2. ] ps,
 'coordinates': [[[-0.72188784  0.74636066  0.05813137]
   [-0.81043531  0.64553013 -0.00369961]
   [-0.73617225  0.5198318  -0.0486595 ]
   ...
   [-0.30431123  0.95345676  0.02049058]
   [-0.36000762  0.93144413  0.18751983]
   [-0.12497457  0.86743548  0.21938995]]
 
  [[-0.72297679  0.75950193  0.1033608 ]
   [-0.81054603  0.6797014   0.011069  ]
   [-0.75025977  0.54347444 -0.02469187]
   ...
   [-0.28984313  0.94931681  0.00860516]
   [-0.34914685  0.93085583  0.18166777]
   [-0.09713665  0.85020737  0.16878855]]
 
  [[-0.69809542  0.79990165  0.11378251]
   [-0.79235712  0.73827248  0.01545496]
   [-0.75864621  0.5912139  -0.01251315]
   ...
   [-0.3001207   0.95915799  0.04391413]
   [-0.33717406  0.91896417  0.21398136]
   [-0.10241378  0.84819644  0.18889462]]
 
  ...
 
  [[-0.62782934  0.85382559 -0.00604494]
   [-0.67765962  0.84164945 -0.14239682]
   [-0.69998651  0.69796684 -0.170781  ]
   ...
   [-0.22086467  0.93736667  0.13933774]
   [-0.2722364   0.85569963  0.28860563]
   [-0.04800618  0.78998921  0.31050751]]
 
  [[-0.60239314  0.86971734 -0.00639182]
   [-0.6600158   0.84148465 -0.13896391]
   [-0.68568053  0.69166892 -0.16360896]
   ...
   [-0.18681713  0.95682869  0.04510207]
   [-0.3072922   0.92063274  0.17565054]
   [-0.04693417  0.85803783  0.23513554]]
 
  [[-0.65077333  0.85254932 -0.00292442]
   [-0.7131694   0.82200085 -0.13299924]
   [-0.71638328  0.66963951 -0.1591624 ]
   ...
   [-0.21179664  0.96669136 -0.00634201]
   [-0.23525693  0.93727469  0.17397302]
   [-0.01641677  0.84405769  0.14259919]]] nm,
 'velocities': [[[ 4.81170017e-01  4.76366786e-01  4.40211556e-01]
   [ 6.83519358e-02  5.39074700e-01 -1.03908109e+00]
   [ 2.93295337e-01  2.61587578e-01  5.86006096e-01]
   ...
   [-2.04255932e+00 -4.84470024e-01  1.04722546e-01]
   [-2.85259010e-01  3.23113544e+00 -1.67397751e+00]
   [ 8.37677414e-01  2.57887275e+00  2.64165984e+00]]
 
  [[ 1.50162744e-01 -1.14965086e-01 -2.70265252e-01]
   [ 4.78291376e-01  7.22742156e-01 -2.70259747e-01]
   [-5.55101101e-01  6.17024107e-01  1.06677396e-01]
   ...
   [-3.27096365e-01  4.68811108e-01 -3.53511782e-01]
   [-1.26986183e+00  8.22781674e-01 -1.22555843e+00]
   [-1.03503384e-01 -6.42982198e-01  1.06130673e-01]]
 
  [[ 1.30296262e-01 -1.30475548e-01  2.88456585e-01]
   [ 2.73391076e-01  4.63430436e-01 -2.27647287e-01]
   [-4.00074360e-02  1.32426338e-01 -5.13497873e-01]
   ...
   [-9.84989232e-01 -5.44261722e-01  3.34381846e-01]
   [ 1.95064436e-02 -1.99996258e+00  7.21349167e-01]
   [-8.40840476e-01 -6.43519603e-01  5.50726734e-01]]
 
  ...
 
  [[ 5.95113119e-01  6.58859467e-01  1.06488871e-01]
   [-4.07145068e-01  2.60519416e-03  4.55723216e-02]
   [ 7.76822778e-01 -5.01100602e-01 -2.40298476e-01]
   ...
   [ 8.08932631e-01  8.01310795e-02 -9.82690441e-01]
   [ 1.90905995e+00  7.35494064e-01  1.40260125e+00]
   [-7.13687953e-01 -8.23588295e-01 -4.60359886e-01]]
 
  [[-6.77492552e-01  3.62104507e-01  9.97444072e-02]
   [ 4.14625977e-01  3.25770200e-02 -1.55284970e-01]
   [-3.15142757e-01 -9.21834143e-01 -7.60452446e-01]
   ...
   [ 1.34800106e+00 -9.02971007e-01 -1.16291450e+00]
   [-8.12642810e-01  3.26561134e-01 -1.96930429e+00]
   [ 2.61510243e+00  2.25010744e+00  2.33677387e+00]]
 
  [[ 4.01109404e-01 -1.41797581e-01  1.06423344e-01]
   [ 2.49801564e-01  2.16378316e-01 -1.27145017e-01]
   [ 7.42880295e-01 -4.80956527e-01  2.85312267e-01]
   ...
   [ 2.66513095e+00 -1.69046724e+00 -3.57891098e+00]
   [ 2.57081089e+00  2.13434363e-01  5.54696885e-01]
   [ 5.22573347e-01 -2.44720460e-02 -4.51785075e-01]]] nm/ps,
 'potential_energy': [-2122.86180618 -1458.03823061 -1467.11853366 -1419.16870708
  -1390.68920689 -1291.91272309 -1304.44485124 -1202.79498215
  -1143.68475883 -1136.51280433 -1160.36936942] kJ/mol,
 'kinetic_energy': [1050.17880481  586.10721076  676.55973727  727.11400407  764.28199896
   747.62714695  810.30324902  839.20860611  907.83361075  909.77905585
   953.16677122] kJ/mol,
 'total_energy': [-1072.68300137  -871.93101985  -790.55879639  -692.05470301
   -626.40720793  -544.28557614  -494.14160223  -363.58637604
   -235.85114809  -226.73374848  -207.20259821] kJ/mol,
 'temperature': [332.82602917 185.75097377 214.41747807 230.43929817 242.21869811
  236.94038909 256.80390002 265.96467834 287.71353452 288.33009121
  302.08066488] K,
 'box': [[[2. 0. 0.]
   [0. 2. 0.]
   [0. 0. 2.]]
 
  [[2. 0. 0.]
   [0. 2. 0.]
   [0. 0. 2.]]
 
  [[2. 0. 0.]
   [0. 2. 0.]
   [0. 0. 2.]]
 
  [[2. 0. 0.]
   [0. 2. 0.]
   [0. 0. 2.]]
 
  [[2. 0. 0.]
   [0. 2. 0.]
   [0. 0. 2.]]
 
  [[2. 0. 0.]
   [0. 2. 0.]
   [0. 0. 2.]]
 
  [[2. 0. 0.]
   [0. 2. 0.]
   [0. 0. 2.]]
 
  [[2. 0. 0.]
   [0. 2. 0.]
   [0. 0. 2.]]
 
  [[2. 0. 0.]
   [0. 2. 0.]
   [0. 0. 2.]]
 
  [[2. 0. 0.]
   [0. 2. 0.]
   [0. 0. 2.]]
 
  [[2. 0. 0.]
   [0. 2. 0.]
   [0. 0. 2.]]] nm}