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

# # Replacing parameters in Interchange
# 
# <details>
#     <summary><small>▼ Click here for dependency installation instructions</small></summary>
#     The simplest way to install dependencies is to use the Interchange examples environment. From the root of the cloned openff-interchange repository:
#     
#     conda env create --name interchange-examples --file devtools/conda-envs/examples_env.yaml
#     conda activate interchange-examples
#     pip install -e .
#     cd examples/parameter_replacement
#     jupyter notebook parameter_replacement.ipynb
#     
# </details>
# 
# In this example, we'll compare the parametrization of ethanol in the Parsley and Sage force fields, and then update the former system with Sage angle terms. This is an advanced example for users looking to optimize force fields.

# In[ ]:


from openff.toolkit import ForceField, Molecule

from openff.interchange import Interchange


# Load in the OpenFF force fields Parsley and Sage. These are the first two major versions of the OpenFF force field family. They are distributed as part of the `openff-forcefields` conda package, which is a dependency of the Toolkit.

# In[ ]:


pars_ff = ForceField("openff-1.3.1.offxml")
sage_ff = ForceField("openff-2.0.0.offxml")


# Now, we'll construct Interchanges describing a system of two ethanol molecules for each force field.

# In[ ]:


ethanol = Molecule.from_smiles("CCO")

pars_sys = Interchange.from_smirnoff(force_field=pars_ff, topology=[ethanol] * 2)
sage_sys = Interchange.from_smirnoff(force_field=sage_ff, topology=[ethanol] * 2)


# ## Inspecting parameters

# Sage changes quite a lot from Parsley (see the [Release Notes]), but we'll focus on the bond angle parameters here. Let's first confirm that the angle parameters have changed. Parameters in Interchange are managed by [`PotentialHandler`] objects; each kind of potential has its own subclass. Angles in SMIRNOFF force fields use [`SMIRNOFFAngleHandler`]. Potential handlers are similar to parameter handlers, which fill a similar role in force fields, but unlike parameter handlers potential handlers connect the topology to the potential. Parameter handlers merely connect SMIRKS codes to parameters. For more information, see [Tweaking and Inspecting Parameters] in the user guide.
# 
# We can get the potential handler for each Interchange from the `handlers` attribute. Let's check that Parsley and Sage have different bond angle potentials by comparing the appropriate handlers.
# 
# [Release Notes]: https://github.com/openforcefield/openff-forcefields/releases
# [`PotentialHandler`]: https://openff-interchange.readthedocs.io/en/stable/_autosummary/openff.interchange.components.potentials.PotentialHandler.html
# [`SMIRNOFFAngleHandler`]: https://openff-interchange.readthedocs.io/en/stable/_autosummary/openff.interchange.components.smirnoff.SMIRNOFFAngleHandler.html
# [Tweaking and Inspecting Parameters]: https://openff-interchange.readthedocs.io/en/stable/using/handlers.html

# In[ ]:


assert pars_sys.handlers["Angles"] != sage_sys.handlers["Angles"]


# Handlers map between topology and potential via two dictionary attributes, `.key_map` and `.potentials`. There's a bit of [complexity] here but essentially `key_map` maps from atom indices to an identifier, while `potentials` maps from the identifier to actual parameters. This means that getting the parameters for a particular atom or group of atoms is as fast as two dict lookups, without duplicating parameters. SMIRNOFF handlers use SMIRKS codes for the identifier, but other handlers can use whatever they want.
# 
# [complexity]: https://openff-interchange.readthedocs.io/en/stable/using/handlers.html
# 
# Parsley and Sage use the same SMIRKS codes for their ethanol angle terms, and both our systems have the same topologies, so their `key_map` attributes should be identical:

# In[ ]:


assert pars_sys.handlers["Angles"].key_map == sage_sys.handlers["Angles"].key_map


# Since they have the same slot maps, they should have the same SMIRKS codes, so the keys to `.potentials` should be the same:

# In[ ]:


assert (
    pars_sys.handlers["Angles"].potentials.keys()
    == sage_sys.handlers["Angles"].potentials.keys()
)


# So we can compare them key-by-key. Here's Parsley:

# In[ ]:


pars_sys.handlers["Angles"].potentials


# And Sage:

# In[ ]:


sage_sys.handlers["Angles"].potentials


# Here's the relevant bits next to each other. See if you can find them in the above!

# In[ ]:


for potential_key in pars_sys.handlers["Angles"].potentials:
    print(potential_key.id)
    for ff, sys in [("Sage", sage_sys), ("Parsley", pars_sys)]:
        print(f"    {ff}")
        for k, v in sys.handlers["Angles"].potentials[potential_key].parameters.items():
            print(f"        {k}: {v}")


# ## Updating Parsley with Sage angles

# <div class="alert alert-danger" style="max-width: 700px; margin-left: auto; margin-right: auto;">
#     <b>🛑 For Educational Purposes Only</b><br />
#     Don't combine parameters from different force fields or "update" one part of a force field without careful consideration. Parameters are designed and optimized to work within a particular context, and moving them to a completely different force field is usually a bad idea.
# </div>
# 
# Just to demonstrate changing parameters, lets "update" our Parsley Interchange to use Sage angle parameters. This will definitely hurt force field performance, as each set of angle parameters are optimized for different non-bonded parameters, but it demonstrates the API, and we won't be running any simulations.

# Generally, updating a `PotentialHandler` involves updating both dictionaries, `key_map` and `potentials`. Since both of our Interchanges use the same SMIRKS codes and parameter types, we could skip updating the `key_map`, but it's best practice to do both steps every time. It's good practice to clear the old values so that you don't accidentally wind up with redundant parameters (unless that's what you want), so let's start with that:

# In[ ]:


pars_sys.handlers["Angles"].key_map.clear()
pars_sys.handlers["Angles"].potentials.clear()


# The [`store_matches`] method updates `key_map`. It takes a topology and some object that describes the parameters (for SMIRNOFF, a `ParameterHandler`) and applies the described force field, storing the appropriate atom indices in the calling handler. If there are clashes, the new values replace the old.
# 
# [`store_matches`]: https://openff-interchange.readthedocs.io/en/stable/_autosummary/openff.interchange.components.potentials.PotentialHandler.html#openff.interchange.components.potentials.PotentialHandler.store_matches

# In[ ]:


pars_sys.handlers["Angles"].store_matches(sage_ff["Angles"], topology=pars_sys.topology)


# The [`store_potentials`] method updates `potentials`. It just takes the object that describes the parameters and stores them. If there are clashes, the new values replace the old.
# 
# [`store_potentials`]: https://openff-interchange.readthedocs.io/en/stable/_autosummary/openff.interchange.components.potentials.PotentialHandler.html#openff.interchange.components.potentials.PotentialHandler.store_potentials

# In[ ]:


pars_sys.handlers["Angles"].store_potentials(sage_ff["Angles"])


# ### Check the update

# We can check that the two systems now have identical angle handlers, and that other parameters, say the bonds, are still the same:

# In[ ]:


assert pars_sys.handlers["Angles"] == sage_sys.handlers["Angles"]
assert pars_sys.handlers["Bonds"] != sage_sys.handlers["Bonds"]


# <!-- TODO: Round trip to disk -->