from openmm.app import *
from openmm import *
from openmm.unit import picosecond, second, nanometer, kelvin, millisecond
from sys import stdout
import time

# Load in the PDB strucure
pdb = PDBFile("pdbs/villin.pdb")

# Specifiy the forcefield
forcefield = ForceField("amber14-all.xml", "amber14/tip3pfb.xml")


# Combine the molecular topology and the forcefield
system = forcefield.createSystem(
    pdb.topology, nonbondedMethod=PME, nonbondedCutoff=1 * nanometer, constraints=HBonds
)

time_step = 0.004 * picosecond
total_time = 50 * picosecond
steps = int(total_time / time_step)

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

platform = Platform.getPlatform("OpenCL")
simulation = Simulation(pdb.topology, system, integrator, platform)
simulation.context.setPositions(pdb.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
    )
)

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