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

# # Methods for Topology solvation

# This example explores some different ways to solvate an OpenFF Topology. See also the "Generating and Parametrizing multi-component systems", "Solvating and equilibrating a ligand in a box of water", and the Toolkit Showcase.

# In[ ]:


from typing import Optional

import nglview
import numpy
import openmm
import openmm.app
import openmm.unit
from openff.toolkit import ForceField, Molecule, Topology
from openff.units import Quantity, unit
from openff.units.openmm import ensure_quantity
from openmm.app import Topology as OpenMMTopology

from openff.interchange import Interchange


# ## Methods

# In[ ]:


OPENMM_IONS = {
    "Li+": "[#3+1]",
    "Na+": "[#11+1]",
    "K+": "[#19+1]",
    "Rb+": "[#37+1]",
    "Cs+": "[#55+1]",
    "F-": "[#9-1]",
    "Cl-": "[#17-1]",
    "Br-": "[#35-1]",
    "I-": "[#53-1]",
}


def visualize_all(interchange: Interchange) -> nglview.NGLWidget:
    view = interchange.visualize()
    view.clear_representations()
    view.add_representation("ball+stick", selection="all")
    return view


def solvate_topology(
    topology: Topology,
    method: str = "pdbfixer",
    box_vectors: Optional[Quantity] = Quantity(5.0 * numpy.ones(3), unit.nanometer),
    **kwargs,
) -> Topology:
    if method in ["pdbfixer", "openmm"]:
        boxSize = openmm.unit.Quantity(
            openmm.Vec3(*box_vectors.m_as(unit.nanometer)), openmm.unit.nanometer
        )

        if method == "pdbfixer":
            openmm_topology, openmm_positions = _solvate_pdbfixer(
                topology.to_openmm(),
                topology.get_positions().to_openmm(),
                boxSize=boxSize,
                **kwargs,
            )
        else:
            openmm_topology, openmm_positions = _solvate_openmm(
                topology.to_openmm(),
                topology.get_positions().to_openmm(),
                boxSize=boxSize,
                **kwargs,
            )

        unique_molecules: List[Molecule] = [*topology.unique_molecules]
        unique_molecules.append(Molecule.from_mapped_smiles("[H:2][O:1][H:3]"))

        if "positiveIon" in kwargs:
            unique_molecules.append(
                Molecule.from_smiles(OPENMM_IONS[kwargs["positiveIon"]])
            )

        if "negativeIon" in kwargs:
            unique_molecules.append(
                Molecule.from_smiles(OPENMM_IONS[kwargs["negativeIon"]])
            )

        new_topology = Topology.from_openmm(
            openmm_topology,
            unique_molecules=unique_molecules,
        )

        new_topology.set_positions(ensure_quantity(openmm_positions, "openff"))

        return new_topology


def _solvate_pdbfixer(
    topology: OpenMMTopology,
    positions: openmm.unit.Quantity,
    **kwargs,
) -> tuple[OpenMMTopology, openmm.unit.Quantity]:
    """
    Add solvent and ions using PDBFixer.

    https://htmlpreview.github.io/?https://github.com/openmm/pdbfixer/blob/master/Manual.html

    """
    import pdbfixer

    with open("_tmp.pdb", "w") as _file:
        openmm.app.PDBFile.writeFile(topology, positions, _file)

    pdb_object = pdbfixer.PDBFixer("_tmp.pdb")
    pdb_object.addSolvent(**kwargs)

    return pdb_object.topology, pdb_object.positions


def _solvate_openmm(
    topology: OpenMMTopology,
    positions: openmm.unit.Quantity,
    box_vectors: openmm.unit.Quantity,
    forcefield: Optional[openmm.app.ForceField] = None,
    **kwargs,
) -> tuple[OpenMMTopology, openmm.unit.Quantity]:
    if not forcefield:
        import pdbfixer

        forcefield = pdbfixer.PDBFixer._createForceField(topology)

    modeller = openmm.app.Modeller(topology, positions)
    modeller.addSolvent(
        forcefield,
        **kwargs,
    )


# ## Solvating and Parametrizing

# Load the force field:

# In[ ]:


sage_ff14sb = ForceField("openff-2.0.0.offxml", "ff14sb_off_impropers_0.0.3.offxml")


# Load the solute from a PDB file:

# In[ ]:


peptide = Molecule.from_polymer_pdb("ace-a5ca5-nme.pdb")


# Solvate:

# In[ ]:


solvated_topology = solvate_topology(
    peptide.to_topology(),
    box_vectors=Quantity(5.0 * numpy.ones(3), unit.nanometer),
)


# Parametrize and visualize:

# In[ ]:


interchange = Interchange.from_smirnoff(sage_ff14sb, solvated_topology)
visualize_all(interchange)


# And finally, some alternative solvation strategies:

# In[ ]:


solvated_topology = solvate_topology(
    peptide.to_topology(),
    box_vectors=Quantity([10, 4, 4], unit.nanometer),
)


# In[ ]:


solvated_topology = solvate_topology(
    peptide.to_topology(),
    box_vectors=Quantity(5.0 * numpy.ones(3), unit.nanometer),
    positiveIon="K+",
    negativeIon="Br-",
    ionicStrength=2.0 * openmm.unit.molar,
)