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

#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)

tqdm_reporter = msm.thirds.openmm.reporters.H5MSMReporter('test.h5msm', 100, 1000)

simulation.reporters.append(tqdm_reporter)

simulation.step(1000)

tqdm_reporter.close()

molsys = msm.convert('test.h5msm')

msm.info(molsys)

form	n_atoms	n_groups	n_components	n_chains	n_molecules	n_entities	n_peptides	n_structures
molsysmt.MolSys	304	20	1	1	1	1	1	11

os.remove('test.h5msm')