#!/usr/bin/env python
# coding: utf-8

# # Load a system from OpenMM into Interchange

# Interchange is intended as a bidirectional system conversion tool, but it's functionality for loading systems from software other than the OpenFF stack is experimental and in active development. This example describes creating an `Interchange` from an OpenMM `System` and `Topology`.

# In[ ]:


import pathlib
import urllib.request

import openmm
import openmm.app
import openmm.unit
from openff.units.openmm import ensure_quantity

from openff.interchange import Interchange
from openff.interchange.drivers import get_openmm_energies
from openff.interchange.drivers.openmm import _get_openmm_energies, _process


# First we'll create an OpenMM system and topology to experiment on:

# In[ ]:


def openmm_pathway(
    pdb_file: openmm.app.PDBFile,
    force_field: openmm.app.ForceField,
) -> tuple[
    openmm.System,
    openmm.unit.Quantity,
    openmm.unit.Quantity,
]:
    """
    Return an openmm.System, coordinates, and box vectors that result from applying this force field
    to this PDB file.
    """
    from openmm import unit

    system = force_field.createSystem(
        pdb_file.topology,
        nonbondedCutoff=unit.Quantity(0.9, unit.nanometer),
    )

    return system, pdb_file.getPositions(), pdb_file.topology.getPeriodicBoxVectors()


# In[ ]:


if not pathlib.Path("protein.pdb").is_file():
    urllib.request.urlretrieve(
        "https://raw.githubusercontent.com/bayer-science-for-a-better-life/abfe-benchmark/a3b15632d2f419857e2ba73a6922de6f09a66caa/structures/brd4/protein.pdb",
        "brd4.pdb",
    )

protein = openmm.app.PDBFile("brd4.pdb")
amber_force_fields = openmm.app.ForceField("amber14-all.xml", "amber14/tip3pfb.xml")

openmm_system = amber_force_fields.createSystem(
    topology=protein.topology,
    nonbondedMethod=openmm.app.PME,
    nonbondedCutoff=9.0 * openmm.unit.angstrom,
    switchDistance=8.0 * openmm.unit.angstrom,
)


# Then, load it into Interchange. Note that the topology mentioned here is an OpenMM topology, **not** an OpenFF one.

# In[ ]:


converted_interchange = Interchange.from_openmm(
    topology=protein.topology,
    system=openmm_system,
    positions=ensure_quantity(protein.getPositions(), "openff"),
    box_vectors=ensure_quantity(protein.topology.getPeriodicBoxVectors(), "openff"),
)


# As a partial validation of `Interchange.from_openmm`, we can compute the single-point energy of this configuration as generated by each code path. We use a high-level `get_openmm_energies` function that takes in only an `Interchange` object to get the energy of `converted_interchange` and a private `_get_openmm_energies` function to process the OpenMM objects generated by the OpenMM code path.

# In[ ]:


converted_energies = get_openmm_energies(converted_interchange, combine_nonbonded_forces=True)
print(converted_energies)


# In[ ]:


(
    reference_system,
    reference_positions,
    reference_box_vectors,
) = openmm_pathway(protein, amber_force_fields)

reference_energy = _process(
    _get_openmm_energies(
        system=reference_system,
        positions=reference_positions,
        box_vectors=reference_box_vectors,
        round_positions=None,
        platform="Reference",
    ),
    system=reference_system,
    combine_nonbonded_forces=True,
    detailed=False,
)

print(reference_energy)


# The non-bonded energies differ slightly. It's not currently clear if this is due to difference in non-bonded settings or a subtle difference in the conversion of applied parameters. Energies from valence terms match to a high level of precision.