import os
from numpy import diag
from simtk.openmm.app import *
from efieldreporter_h5 import EFieldReporter
from simtk.openmm import *
from simtk.unit import *

temperature=300*kelvin
cuda=openmm.Platform_getPlatform(1)
reference=openmm.Platform_getPlatform(0)

Input='Frame7784-State0'
Out=Input+"/"
PDBFilename=Input+".pdb"

os.mkdir(Out)

pdb = PDBFile(PDBFilename)

forcefield = ForceField('amber10.xml',"tip3p.xml","PCN.xml")

system = forcefield.createSystem(pdb.topology, nonbondedMethod=PME,nonbondedCutoff=1*nanometer, constraints=HBonds)
system.addForce(AndersenThermostat(300*kelvin, 1/picosecond))
system.addForce(MonteCarloBarostat(1*bar, 300*kelvin))

integrator = VerletIntegrator(0.002*picoseconds)

simulation = Simulation(pdb.topology, system, integrator)
simulation.context.setPositions(pdb.positions)

integrator2 = VerletIntegrator(0.002*picoseconds)
system2 = forcefield.createSystem(pdb.topology, nonbondedMethod=PME,nonbondedCutoff=1*nanometer, constraints=HBonds)
simulation2 = Simulation(pdb.topology, system2, integrator2,platform=reference)

integrator3 = VerletIntegrator(0.002*picoseconds)
system3 = forcefield.createSystem(pdb.topology, nonbondedMethod=PME,nonbondedCutoff=1*nanometer, constraints=HBonds)
simulation3 = Simulation(pdb.topology, system3, integrator3,platform=reference)

simulation.reporters.append(EFieldReporter(Out+'a150.h5', 100,simulation2,150))
simulation.reporters.append(EFieldReporter(Out+'a151.h5', 100,simulation3,151))
simulation.reporters.append(EFieldReporter(Out+'a254.h5', 100,simulation2,254))
simulation.reporters.append(EFieldReporter(Out+'a255.h5', 100,simulation3,255))
simulation.reporters.append(PDBReporter(Out+"production.pdb", 5000))

simulation.step(2500000)