# Input Files
print('Parsing PDB file...')
pdb = PDBFile('input/ala2.pdb')

print('Parsing FF file(s)...')
forcefield = ForceField('amber99sbildn.xml')


"""
Explanation: Let us run a MD simulation of alanine dipeptide in vacuum using the AMBER 99 forcefield: first we need an input object (e.g. a pdb), and a forcefield definition (several are provided with the OpenMM installation)
End of explanation
"""


# System Configuration
nonbondedMethod = CutoffNonPeriodic
nonbondedCutoff = 1.4*nanometers
constraints = HBonds
constraintTolerance = 0.00001

"""
Explanation: Define simulation parameters here: periodic boundaries if any, nonbonded interaction cutting scheme, rigid constraints for some of the bonds, etc.
End of explanation
"""


# Integration Options
dt = 0.002*picoseconds
temperature = 300.00*kelvin
friction = 2/picosecond

integrator = LangevinIntegrator(temperature, friction, dt)
integrator.setConstraintTolerance(constraintTolerance)


"""
Explanation: Create an Integrator : we use 'Langevin' at a temperature of 300 K ; for systems where a simulation box is defined it is also possible to add a barostat and then run in the NPT ensemble
End of explanation
"""


platform = Platform.getPlatformByName('CPU')
platformProperties = {'CpuThreads':'1'}

#platform = Platform.getPlatformByName('OpenCL')
#platformProperties = {'OpenCLPrecision':'mixed'}

#platform = Platform.getPlatformByName('CUDA')
#platformProperties = {'CudaPrecision': 'mixed'}


"""
Explanation: An OpenMM platform object is used for providing a level of hardware abstraction : one can choose Reference, CPU, OpenCL or CUDA
End of explanation
"""


# Prepare the Simulation
print('Building system...')
topology = pdb.topology
positions = pdb.positions
system = forcefield.createSystem(topology, nonbondedMethod=nonbondedMethod,
                                 nonbondedCutoff=nonbondedCutoff,
                                 constraints=constraints)

simulation = Simulation(topology, system, integrator,
                        platform, platformProperties)

# set initial positions to what was one obtained from PDB parsing
simulation.context.setPositions(positions)


"""
Explanation: A system object is created from the pdb, the forcefield and some of the simulation parameters above defined : it contains coordinates, velocities,
atom types, and all the data required for energy/force evaluations.
But the computation is performed within a simulation object (using the above defined platform which actually contains the harware dependent compiled code).
End of explanation
"""


# Minimize energy
print('Performing energy minimization...')
simulation.minimizeEnergy()

# after minimization of the energy, set initial velocities to random distribution for a given temperature
simulation.context.setVelocitiesToTemperature(temperature)


"""
Explanation: Energy minimization is performed using the L-BFGS method (https://en.wikipedia.org/wiki/Limited-memory_BFGS)
End of explanation
"""


equilibrationSteps = 25000
printFreq = 250

# reporters objects are used for terminal IO and trajectory IO
#  here we use the StateDataReporter for regularly printing simulation time, energy, etc. to a terminal
simulation.reporters.append(StateDataReporter(stdout, printFreq, step=True, 
    time=True, potentialEnergy=True, kineticEnergy=True, temperature=True))

print('Equilibrating...')
simulation.step(equilibrationSteps)

"""
Explanation: Perform some steps of temperature equilibration :
End of explanation
"""


# reste step and time counters
simulation.currentStep = 0
simulation.context.setTime(0.0)

# save the equilibration results to file : state is platform independent but less precise, checkpoint file
#  contains the binary representation but it is platform dependent
simulation.saveState('output/eq.state')
simulation.saveCheckpoint('output/eq.chk')

simsteps = 5e5 # 1 ns
svfreq=500

simulation.reporters = []

simulation.reporters.append(StateDataReporter(stdout, svfreq, step=True, 
    time=True, potentialEnergy=True, kineticEnergy=True, totalEnergy=True,
    temperature=True, progress=True,remainingTime=True, speed=True, totalSteps=simsteps))

hdf5rep = mdt.reporters.HDF5Reporter('output/sim_1ns.hdf5', svfreq,
                                     coordinates=True, time=True, cell=True, potentialEnergy=True,
                                     kineticEnergy=True, temperature=True, velocities=True)

# dcdrep = DCDReporter('output/sim_1ns.dcd', svfreq)

simulation.reporters.append(hdf5rep)
# simulation.reporters.append(dcdrep)

simulation.step(simsteps)

simulation.saveState('output/sim_1ns.state')
simulation.saveCheckpoint('output/sim_1ns.chk')

hdf5rep.close()


"""
Explanation: Now that the system is temperature equilirated we can run the simulation ; we use the mdtraj package which provides extra trajectory formats (e.g. HDF5)
End of explanation
"""


traj = mdt.load_hdf5("output/sim_1ns.hdf5")

view = ng.show_mdtraj(traj)
view.add_ball_and_stick("all")

view

"""
Explanation: When simulation finished we can have a look at the trajectory using an external trajectory viewer (PyMOL, VMD,...) or using nglview within this jupyter notebook:
End of explanation
"""


