from simtk.openmm.app import PDBFile
from openff.toolkit.utils import get_data_file_path

"""
Explanation: Running a simulation of a System parametrized with a SMIRNOFF force field
This example shows how to use the Open Force Field Toolkit to create a parametrized System object that can be used to run a molecular dynamic simulation with OpenMM. If you want to run MD with a different engine, see the example in examples/conversion_amber_gromacs/.
Create an OpenMM System
We start by loading a PDB file containing one copy of ethanol and cyclohexane. Our goal is to create an OFF Topology object describing this system that we can parametrize with the SMIRNOFF-format "Parsley" force field.
The two Molecule objects created from the SMILES strings can contain information such as partial charges and stereochemistry that is not included in an OpenMM topology. In this example, partial charges are not explicitly given, and ForceField will assign AM1/BCC charges as specified by the "Parsley" force field.
End of explanation
"""


# 1 molecule of ethanol and 1 of cyclohexane.
pdb_file_path = get_data_file_path('systems/test_systems/1_cyclohexane_1_ethanol.pdb')

# 40%-60% cyclohexane-ethanol mixture.
# pdb_file_path = get_data_file_path('systems/packmol_boxes/cyclohexane_ethanol_0.4_0.6.pdb')

pdbfile = PDBFile(pdb_file_path)

"""
Explanation: We use the get_data_file_path utility function to easily access the data installed with the toolkit. Here you have the option to load example systems of increasing complexity. For speed, we recommend that you begin by loading a system with a single ethanol and a single cyclohexane.
End of explanation
"""


from openff.toolkit.topology import Molecule
ethanol = Molecule.from_smiles('CCO')
cyclohexane = Molecule.from_smiles('C1CCCCC1')

"""
Explanation: PDB files are not a reliable source of bond orders, so the toolkit requires users to supply a more detailed description of the molecule and its connectivity (currently either a SMILES, sdf, or mol2 file).
End of explanation
"""


from openff.toolkit.utils.toolkits import OpenEyeToolkitWrapper
if OpenEyeToolkitWrapper.is_available():
    ethanol = Molecule.from_file(get_data_file_path('molecules/ethanol.mol2'))
    cyclohexane = Molecule.from_file(get_data_file_path('molecules/cyclohexane.mol2'))

"""
Explanation: Alternatively, if you have sdf files of the molecules, or if you have OpenEye installed and mol2 files available, you can get the same results as above by loading the detailed molecule information from the files.
End of explanation
"""


from openff.toolkit.topology import Topology
from openff.toolkit.typing.engines.smirnoff import ForceField

# Create the Open Force Field Toolkit Topology from an OpenMM Topology object.
omm_topology = pdbfile.topology
off_topology = Topology.from_openmm(omm_topology, unique_molecules=[ethanol, cyclohexane])

# Load the OpenFF "Parsley" force field.
forcefield = ForceField('openff-1.0.0.offxml')

# Parametrize the topology and create an OpenMM System.
system = forcefield.create_openmm_system(off_topology)

"""
Explanation: We now create the Open Force Field Toolkit Topology describing the system from an OpenMM Topology object. The OFF Topology include more information (supplied by the two Molecule objects) than the OpenMM Topology such as (optionally) partial charges and stereochemistry. In this example, partial charges are not explicitly given, and ForceField will assign AM1/BCC charges as specified by the "Parsley" force field.
Note that the Open Force Field Toolkit produces deterministic charges that do not depend on the input conformation of parameterized molecules. See the FAQ for more information.
<div class="alert alert-block alert-info">
    <b>Note on partial charges:</b> The full 1.0.0 release will implement support for the definition of semiempirical partial charge treatment directly into the SMIRNOFF force field file (for details, see https://open-forcefield-toolkit.readthedocs.io/en/latest/smirnoff.html#partial-charge-and-electrostatics-models). Moreover, it will be possible to import charges directly from sdf and mol2 files.
</div>
End of explanation
"""


from simtk import openmm, unit

# Propagate the System with Langevin dynamics.
time_step = 2*unit.femtoseconds  # simulation timestep
temperature = 300*unit.kelvin  # simulation temperature
friction = 1/unit.picosecond  # collision rate
integrator = openmm.LangevinIntegrator(temperature, friction, time_step)

# Length of the simulation.
num_steps = 1000  # number of integration steps to run

# Logging options.
trj_freq = 1  # number of steps per written trajectory frame
data_freq = 1  # number of steps per written simulation statistics

# Set up an OpenMM simulation.
simulation = openmm.app.Simulation(omm_topology, system, integrator)

# Set the initial positions.
positions = pdbfile.getPositions() 
simulation.context.setPositions(positions)

# Randomize the velocities from a Boltzmann distribution at a given temperature.
simulation.context.setVelocitiesToTemperature(temperature)

# Configure the information in the output files.
pdb_reporter = openmm.app.PDBReporter('trajectory.pdb', trj_freq)
state_data_reporter = openmm.app.StateDataReporter('data.csv', data_freq, step=True,
                                                   potentialEnergy=True, temperature=True,
                                                   density=True)
simulation.reporters.append(pdb_reporter)
simulation.reporters.append(state_data_reporter)

import time

print("Starting simulation")
start = time.process_time()

# Run the simulation
simulation.step(num_steps)

end = time.process_time()
print("Elapsed time %.2f seconds" % (end-start))
print("Done!")

"""
Explanation: Run a simulation
We can now use the System object to run molecular dynamics trajectories with OpenMM.
End of explanation
"""


