from openmm.app import *
from openmm import *
from openmm import unit
from openmm.unit import (
    picosecond,
    second,
    nanometer,
    kelvin,
    millisecond,
    angstrom,
    atmosphere,
    femtosecond,
)
from sys import stdout
import time

# Specifiy the forcefield
forcefield = ForceField("amber14/tip4pew.xml")

modeller = Modeller(Topology(), [])
modeller.addExtraParticles(forcefield)
modeller.addSolvent(
    forcefield, model="tip4pew", boxSize=Vec3(3.0, 3.0, 3.0) * nanometer
)

# Combine the molecular topology and the forcefield
system = forcefield.createSystem(
    modeller.topology, nonbondedMethod=PME, constraints=HBonds
)
system.addForce(MonteCarloBarostat(10 * atmosphere, 300 * kelvin, 25))

time_step = 2 * femtosecond
total_time = 2 * unit.nanosecond
steps = int(total_time / time_step)

# The parameters set are temperature, friction coefficient, and timestep.
integrator = LangevinMiddleIntegrator(250 * kelvin, 1 / picosecond, time_step)
# integrator.setTemperature(250*kelvin)

platform = Platform.getPlatform("CUDA")
simulation = Simulation(modeller.topology, system, integrator, platform)
simulation.context.setPositions(modeller.positions)

# Perform local energy minimization
# print("Current energy Minimizing energy...")
# simulation.minimizeEnergy(maxIterations=100)
# print("Energy minimized!")

# Write the trajectory to a file called "output.pdb"
simulation.reporters.append(PDBReporter("output.pdb", steps // 100))

# Report infomation to the screen as the simulation runs
simulation.reporters.append(
    StateDataReporter(
        stdout,
        steps // 100,
        step=True,
        potentialEnergy=True,
        temperature=True,
        density=True,
        progress=True,
        totalSteps=steps,
    )
)

start_time = time.time()
print(f"Starting simulation, {steps} steps")
simulation.step(steps)
print(f"Simulation finished in {time.time() - start_time} seconds")