"""
Utilities for testing.

"""

import collections
import copy
import functools
import importlib
import itertools
import os
import pathlib
import pprint
import textwrap
from contextlib import contextmanager

import numpy as np
import openmm
import pytest
from openff.units.openmm import to_openmm
from openmm import unit as openmm_unit

from openff.toolkit import unit
from openff.toolkit.utils import (
    AmberToolsToolkitWrapper,
    OpenEyeToolkitWrapper,
    RDKitToolkitWrapper,
    get_data_file_path,
)

requires_ambertools = pytest.mark.skipif(
    not AmberToolsToolkitWrapper.is_available(),
    reason="Test requires AmberTools",
)
requires_rdkit = pytest.mark.skipif(
    not RDKitToolkitWrapper.is_available(),
    reason="Test requires RDKit",
)
requires_openeye = pytest.mark.skipif(
    not OpenEyeToolkitWrapper.is_available(),
    reason="Test requires OE toolkit",
)
requires_openeye_mol2 = pytest.mark.skipif(
    not OpenEyeToolkitWrapper.is_available(),
    reason="Test requires OE toolkit to read mol2 files",
)


def _get_readme_path() -> pathlib.Path | None:
    """
    Return a path to the README file or None if it cannot be assured to be the toolkit's README.
    """
    if "site-packages" in __file__:
        # This test file is being collected from the installed package, which
        # does not provide the README file.
        # Note that there will likely be a mis-bundled file
        # $CONDA_PREFIX/lib/python3.x/site-packages/README.md, but this is not
        # the toolkit's README file!
        return None

    else:
        return pathlib.Path(__file__).parents[3] / "README.md"


def has_pkg(pkg_name):
    """
    Helper function to generically check if a package is installed. Intended
    to be used to check for optional dependencies.

    Parameters
    ----------
    pkg_name : str
        The name of the package to check the availability of

    Returns
    -------
    pkg_available : bool
        Boolean indicator if the package is available or not

    Examples
    --------
    >>> has_numpy = has_pkg('numpy')
    >>> has_numpy
    True
    >>> has_foo = has_pkg('other_non_installed_pkg')
    >>> has_foo
    False
    """
    try:
        importlib.import_module(pkg_name)
    except ModuleNotFoundError:
        return False
    return True


def requires_pkg(pkg_name, reason=None):
    """
    Helper function to generate a pytest.mark.skipif decorator
    for any package. This allows tests to be skipped if some
    optional dependency is not found.

    Parameters
    ----------
    pkg_name : str
        The name of the package that is required for a test(s)
    reason : str, optional
        Explanation of why the skipped it to be tested

    Returns
    -------
    requires_pkg : _pytest.mark.structures.MarkDecorator
        A pytest decorator that will skip tests if the package is not available
    """
    if not reason:
        reason = f"Package {pkg_name} is required, but was not found."
    requires_pkg = pytest.mark.skipif(not has_pkg(pkg_name), reason=reason)
    return requires_pkg


@contextmanager
def does_not_raise():
    """A helpful context manager to use inplace of a pytest raise statement
    when no exception is expected."""
    yield


def get_amber_file_path(prefix):
    """Get AMBER prmtop and inpcrd test data filepaths.

    Parameters
    ----------
    prefix : str
        The file name without extension of .prmtop and .inpcrd
        files to retrieve from testdata/systems/amber.

    Returns
    -------
    prmtop_filepath : str
        Absolute path to the AMBER prmtop filepath in testdata/systems/amber
    inpcrd_filepath : str
        Absolute path to the AMBER inpcrd filepath in testdata/systems/amber

    """
    prefix = os.path.join("systems", "amber", prefix)
    prmtop_filepath = get_data_file_path(prefix + ".prmtop")
    inpcrd_filepath = get_data_file_path(prefix + ".inpcrd")
    return prmtop_filepath, inpcrd_filepath


def get_packmol_pdb_file_path(prefix="cyclohexane_ethanol_0.4_0.6"):
    """Get PDB filename for a packmol-generated box

    Parameters
    ----------
    prefix : str, optional, default='cyclohexane_ethanol_0.4_0.6'
        The prefix of .pdb file to retrieve from testdata/systems/packmol_boxes

    Returns
    -------
    pdb_filename : str
        Absolute path to the PDB file
    """
    prefix = os.path.join("systems", "packmol_boxes", prefix)
    pdb_filename = get_data_file_path(prefix + ".pdb")
    return pdb_filename


def get_monomer_mol2_file_path(prefix="ethanol"):
    """Get absolute filepath for a mol2 file denoting a small molecule monomer in testdata

    Parameters
    ----------
    prefix : str, optional, default='ethanol'
        The prefix of .mol2 file to retrieve from systems/monomers/

    Returns
    -------
    mol2_filename : str
        Absolute path to the mol2 file
    """
    # TODO: The mol2 files in this folder are not tripos mol2 files. Delete or convert them.
    prefix = os.path.join("systems", "monomers", prefix)
    mol2_filename = get_data_file_path(prefix + ".mol2")
    return mol2_filename


def extract_compressed_molecules(tar_file_name, file_subpaths=None, filter_func=None):
    if (file_subpaths is None) == (filter_func is None):
        raise ValueError("Only one between file_subpaths and filter_func must be specified.")

    # Find the path of the tarfile with respect to the data/molecules/ folder.
    molecules_dir_path = get_data_file_path("molecules")
    tar_file_path = os.path.join(molecules_dir_path, tar_file_name)
    tar_root_dir_name = tar_file_name.split(".")[0]

    # Return value: Paths to the extracted molecules.
    extracted_file_paths = None

    # Handle subpaths search.
    if file_subpaths is not None:
        # We can already check the paths of the extracted files
        # and skipping opening the tarball if not necessary.
        extracted_file_paths = [
            os.path.join(molecules_dir_path, tar_root_dir_name, file_subpath) for file_subpath in file_subpaths
        ]

        # Remove files that we have already extracted.
        # Also, we augument the subpath with its root directory.
        file_subpaths_set = {
            os.path.join(tar_root_dir_name, subpath)
            for subpath, fullpath in zip(file_subpaths, extracted_file_paths)
            if not os.path.isfile(fullpath)
        }

        # If everything was already extracted, we don't need to open the tarball.
        if len(file_subpaths_set) == 0:
            return extracted_file_paths

        # Otherwise, create a filter matching only the subpaths.
        def filter_func(x):
            return x in file_subpaths_set

    # If no filter was specified, just create one matching everything.
    if filter_func is None:

        def filter_func(x):
            return True

    # Determine opening mode.
    if ".gz" in tar_file_name:
        mode = "r:gz"
    else:
        mode = "r"

    # Extract required files.
    import tarfile

    with tarfile.open(tar_file_path, mode) as tar_file:
        # Gather the paths to extract. Remove the
        members = [m for m in tar_file.getmembers() if filter_func(m.name)]

        # Built the paths to the extracted molecules we didn't already.
        if extracted_file_paths is None:
            extracted_file_paths = [os.path.join(molecules_dir_path, m.name) for m in members]

        # Extract only the members that we haven't already extracted.
        members = [member for member, fullpath in zip(members, extracted_file_paths) if not os.path.isfile(fullpath)]
        tar_file.extractall(path=molecules_dir_path, members=members)

    return extracted_file_paths


def get_alkethoh_file_path(alkethoh_name, get_amber=False):
    """Retrieve the mol2, top and crd files of a molecule in the AlkEthOH set.

    Parameters
    ----------
    alkethoh_name : str
        The name of the AlkEthOH molecule (e.g. "AlkEthOH_r0", "AlkEthOH_c1266").
    get_amber : bool, optional
        If True, the paths to the top and crd files are returned.

    Returns
    -------
    molecule_file_paths : str or list[str]
        All the requested paths. If ``get_amber`` is False, only a single string
        pointing to the path of the mol2 file is returned, otherwise this is a
        list ``[mol2_path, top_path, crd_path]``.

    """
    # Determine if this is a ring or a chain molecule and the subfolder name.
    is_ring = alkethoh_name[9] == "r"
    alkethoh_subdir_name = "rings" if is_ring else "chain"
    alkethoh_subdir_name = "AlkEthOH_" + alkethoh_subdir_name + "_filt1"

    # Determine which paths have to be returned.
    file_base_subpath = os.path.join(alkethoh_subdir_name, alkethoh_name)
    # We always return the mol2 file.
    file_subpaths = [file_base_subpath + "_tripos.mol2"]
    # Check if we need to return also Amber files.
    if get_amber:
        file_subpaths.append(file_base_subpath + ".top")
        file_subpaths.append(file_base_subpath + ".crd")

    return extract_compressed_molecules("AlkEthOH_tripos.tar.gz", file_subpaths=file_subpaths)


def get_freesolv_file_path(freesolv_id, ff_version):
    file_base_name = "mobley_" + freesolv_id
    mol2_file_subpath = os.path.join("mol2files_sybyl", file_base_name + ".mol2")
    xml_dir = "xml_" + ff_version.replace(".", "_")
    xml_file_subpath = os.path.join(xml_dir, file_base_name + "_vacuum.xml")

    # Extract the files if needed.
    file_subpaths = [mol2_file_subpath, xml_file_subpath]
    return extract_compressed_molecules("FreeSolv.tar.gz", file_subpaths=file_subpaths)


def create_system_from_amber(prmtop_file_path, inpcrd_file_path, *args, **kwargs):
    """Create an OpenMM System and Topology from the AMBER files.

    Parameters
    ----------
    prmtop_file_path : str
        Path to the topology/parameter file in AMBER prmtop format.
    inpcrd_file_path : str
        Path to the coordinates file in AMBER inpcrd or rst7 format.
    *args
    **kwargs
        Other parameters to pass to ``openmm.app.AmberPrmtopFile.createSystem``.

    Returns
    -------
    system : openmm.System
        The OpenMM ``System`` object with the parameters.
    topology : openmm.app.Topology
        The OpenMM ``Topology`` object loaded from the AMBER files.
    positions : openmm.unit.Quantity
        Initial positions loaded from the inpcrd or restart file.

    """
    prmtop_file = openmm.app.AmberPrmtopFile(prmtop_file_path)
    # AmberInpcrdFile parses also rst7 files.
    inpcrd_file = openmm.app.AmberInpcrdFile(inpcrd_file_path)
    # Create system and update box vectors (if needed)
    system = prmtop_file.createSystem(*args, **kwargs)
    if inpcrd_file.boxVectors is not None:
        system.setDefaultPeriodicBoxVectors(*inpcrd_file.boxVectors)
    # Store numpy positions.
    positions = inpcrd_file.getPositions(asNumpy=True)
    return system, prmtop_file.topology, positions


def quantities_allclose(quantity1, quantity2, **kwargs):
    """Check if two Quantity objects are close.

    If the quantities are arrays, all their elements must be close.
    If the given arguments are unitless, this simply calls numpy.allclose().

    Parameters
    ----------
    quantity1 : openmm.unit.Quantity
        The first unit to compare.
    quantity2 : openmm.unit.Quantity
        The second unit to compare.
    **kwargs
        Other parameters passed to ``numpy.allclose()``.

    Returns
    -------
    is_close : bool
        True if the two quantities are close, False otherwise.

    """
    if isinstance(quantity1, openmm_unit.Quantity):
        # Check that the two Quantities have compatible units.
        if not quantity1.unit.is_compatible(quantity2.unit):
            raise ValueError(f"The two quantities don't have compatible units: {quantity1.unit} and {quantity2.unit}")
        # Compare the values stripped of the units.
        quantity1 = quantity1.value_in_unit_system(openmm_unit.md_unit_system)
        quantity2 = quantity2.value_in_unit_system(openmm_unit.md_unit_system)
    return np.allclose(quantity1, quantity2, **kwargs)


def get_context_potential_energy(context, positions, box_vectors=None, by_force_group=True):
    """Compute the potential energy of a System in a Context object.

    This is simply a shortcut that takes care of setting
    box vectors and positions before computing the energy.

    Parameters
    ----------
    context : openmm.Context
        The Context object containing the system.
    positions : openmm.unit.Quantity
        A n_atoms x 3 arrays of coordinates with units of length.
    box_vectors : openmm.unit.Quantity, optional
        If specified, this should be a 3 x 3 array. Each row should
        be a box vector.
    by_force_group : bool, optional
        If True, it returns a dictionary mapping force groups to
        their potential energies. Default is True.

    Returns
    -------
    potential_energy : openmm.unit.Quantity or dict[type, Quantity]
        The potential energy of the system at the given coordinates.
        If ``by_force_group`` is True, then this is a dictionary
        mapping force groups to their potential energy.

    """
    # Box vectors must be updated before positions are set.
    if box_vectors is not None:
        context.setPeriodicBoxVectors(*box_vectors)
    context.setPositions(positions)

    # If we don't have energies by force group,
    # just return the energy of the full System.
    if not by_force_group:
        return context.getState(getEnergy=True).getPotentialEnergy()

    # Determine how many force groups we need to compute.
    force_groups = {f.getForceGroup() for f in context.getSystem().getForces()}

    # Compute potential energies of all groups.
    potential_energies = {}
    for force_group in force_groups:
        state = context.getState(getEnergy=True, groups={force_group})
        potential_energies[force_group] = state.getPotentialEnergy()
    return potential_energies


class FailedEnergyComparisonError(AssertionError):
    """Error raised when the energy comparison between two system fails.

    Attributes
    ----------
    potential_energy1
    potential_energy2

    """

    def __init__(self, err_msg, potential_energy1, potential_energy2):
        super().__init__(err_msg)
        self.potential_energy1 = potential_energy1
        self.potential_energy2 = potential_energy2


def compare_context_energies(context1, context2, *args, rtol=1.0e-5, atol=1.0e-8, **kwargs):
    """Compare energies of two Contexts given the same positions.

    Parameters
    ----------
    context1 : openmm.Context
        The first Context object to compare containing the system.
    context2 : openmm.Context
        The second Context object to compare containing the system.
    positions : openmm.unit.Quantity
        A n_atoms x 3 arrays of coordinates with units of length.
    rtol : float, optional, default=1e-5
        The relative tolerance parameter used for the energy comparisons.
    atol : float, optional, default=1e-8
        The absolute tolerance parameter used for the energy comparisons.
    box_vectors : openmm.unit.Quantity, optional
        If specified, this should be a 3 x 3 array. Each row should
        be a box vector.
    by_force_group : bool, optional
        If True, the energies of each force groups are compared.
        Default is True.

    Returns
    -------
    potential_energy1 : openmm.unit.Quantity or dict[int, Quantity]
        The potential energy of context1 at the given coordinates.
        If ``by_force_group`` is True, then this is a dictionary
        mapping force groups to their potential energy.
    potential_energy2 : openmm.unit.Quantity or dict[int, Quantity]
        The potential energy of context2 at the given coordinates.
        If ``by_force_group`` is True, then this is a dictionary
        mapping force groups to their potential energy.

    Raises
    ------
    FailedEnergyComparisonError
        If the potential energies of the two contexts are different
        or if the two Systems in the contexts have forces divided into
        different force groups. The potential energies of context1 and
        context2 can be accessed in the Exception through the attributes
        potential_energy1 and potential_energy2 respectively.

    """
    potential_energy1 = get_context_potential_energy(context1, *args, **kwargs)
    potential_energy2 = get_context_potential_energy(context2, *args, **kwargs)

    def raise_assert(assertion, err_msg, format_args):
        """Shortcut to raise a custom error and format the error message."""
        if not assertion:
            raise FailedEnergyComparisonError(err_msg.format(*format_args), potential_energy1, potential_energy2)

    # If by_force_group is True, then the return value will be a
    # dictionary and we need to compare force group by force group.
    if isinstance(potential_energy1, openmm_unit.Quantity):
        raise_assert(
            assertion=quantities_allclose(potential_energy1, potential_energy2, rtol=rtol, atol=atol),
            err_msg="potential energy 1 {}, potential energy 2: {}",
            format_args=(potential_energy1, potential_energy2),
        )
    else:  # potential_energy1 is a dict.
        # If the two context expose different force groups raise an error.
        raise_assert(
            assertion=set(potential_energy1) == set(potential_energy2),
            err_msg="The two contexts have different force groups: context1 {}, context2 {}",
            format_args=(sorted(potential_energy1), sorted(potential_energy2)),
        )

        # Check the potential energies of all force groups.
        for force_group in potential_energy1:
            energy1 = potential_energy1[force_group]
            energy2 = potential_energy2[force_group]
            raise_assert(
                assertion=quantities_allclose(energy1, energy2, rtol=rtol, atol=atol),
                err_msg="Force group {} do not have the same energies: context1 {}, context2 {}",
                format_args=(force_group, energy1, energy2),
            )

    return potential_energy1, potential_energy2


def compare_system_energies(
    system1,
    system2,
    positions,
    box_vectors=None,
    by_force_type=True,
    ignore_charges=False,
    modify_system=False,
    rtol=1.0e-5,
    atol=1.0e-8,
):
    """Compare energies of two Systems given the same positions.

    Parameters
    ----------
    system1 : openmm.System
        The first Context object to compare containing the system.
    system2 : openmm.System
        The second Context object to compare containing the system.
    positions : openmm.unit.Quantity
        A n_atoms x 3 arrays of coordinates with units of length.
    box_vectors : openmm.unit.Quantity, optional
        If specified, this should be a 3 x 3 array. Each row should
        be a box vector.
    by_force_type : bool, optional
        If True, the contribution to the energy of each force type
        is compared instead of just the overall potential energy.
        Default is True.
    ignore_charges : bool, optional
        If True, particle and exception charges are ignored during
        the comparison. Default is False.
    modify_system : bool, optional
        If either ``by_force_type`` or ``ignore_charges`` are True,
        the function will change the force groups of two ``Systems``
        to compute the separate contributions. Set this to False to
        work on a copy and avoid modifying the original objects.
        Default is False.
    rtol : float, optional, default=1e-5
        The relative tolerance parameter used for the energy comparisons.
    atol : float, optional, default=1e-8
        The absolute tolerance parameter used for the energy comparisons.

    Returns
    -------
    potential_energy1 : openmm.unit.Quantity or dict[str, Quantity]
        The potential energy of context1 at the given coordinates.
        If ``by_force_type`` is True, then this is a dictionary
        mapping force names to their potential energy.
    potential_energy2 : openmm.unit.Quantity or dict[str, Quantity]
        The potential energy of context2 at the given coordinates.
        If ``by_force_type`` is True, then this is a dictionary
        mapping force names to their potential energy.

    Raises
    ------
    FailedEnergyComparisonError
        If the potential energies of the two contexts are different
        or if the two Systems in the contexts have forces divided into
        different force groups. The potential energies of context1 and
        context2 can be accessed in the Exception through the attributes
        potential_energy1 and potential_energy2 respectively.

    """
    # Make a copy of the systems if requested.
    if not modify_system and (by_force_type or ignore_charges):
        system1 = copy.deepcopy(system1)
        system2 = copy.deepcopy(system2)

    # Group force types into OpenMM force groups.
    if by_force_type:
        # Remove eventual CMMotionRemover forces that do not affect energies.
        for system in [system1, system2]:
            for force_idx, force in enumerate(system.getForces()):
                if isinstance(force, openmm.CMMotionRemover):
                    system.removeForce(force_idx)
                    break

        # First check that the two systems have the same force types.
        force_names1 = {f.__class__.__name__ for f in system1.getForces()}
        force_names2 = {f.__class__.__name__ for f in system2.getForces()}
        if set(force_names1) != set(force_names2):
            err_msg = "The two Systems have different forces to compare:\n - system1 {}\n - system2 {}"
            raise ValueError(err_msg.format(sorted(force_names1), sorted(force_names2)))

        # Create a map from force group to force class and viceversa.
        group_to_force = {i: force_name for i, force_name in enumerate(force_names1)}
        force_to_group = {force_name: i for i, force_name in group_to_force.items()}

        # Assign force groups.
        for system in [system1, system2]:
            for force in system.getForces():
                force.setForceGroup(force_to_group[force.__class__.__name__])

    # If we need to ignore the charges, turn off particle and exceptions charges.
    if ignore_charges:
        for system in [system1, system2]:
            # Find the NonbondedForce.
            nonbonded_force = None
            for force in system.getForces():
                if isinstance(force, openmm.NonbondedForce):
                    nonbonded_force = force

            # We checked before that the two systems have the same force
            # types so we can break the loop if there are no NonbondedForces.
            if nonbonded_force is None:
                break

            # Set particle charges to 0.0.
            for particle_idx in range(nonbonded_force.getNumParticles()):
                _, sigma, epsilon = nonbonded_force.getParticleParameters(particle_idx)
                nonbonded_force.setParticleParameters(particle_idx, 0.0, sigma, epsilon)

            # Set particle exceptions charge products to 0.0.
            for exception_idx in range(nonbonded_force.getNumExceptions()):
                (
                    atom1,
                    atom2,
                    _,
                    sigma,
                    epsilon,
                ) = nonbonded_force.getExceptionParameters(exception_idx)
                nonbonded_force.setExceptionParameters(exception_idx, atom1, atom2, 0.0, sigma, epsilon)

    # Create Contexts and compare the energies.
    integrator = openmm.VerletIntegrator(1.0 * openmm_unit.femtoseconds)

    context1 = openmm.Context(system1, integrator, openmm.Platform.getPlatformByName("Reference"))
    context2 = openmm.Context(system2, copy.deepcopy(integrator), openmm.Platform.getPlatformByName("Reference"))

    def map_energies_by_force_type(potential_energy1, potential_energy2):
        """Convert dictionary force_group -> energy to force_type -> energy."""
        potential_energy1 = {group_to_force[group]: energy for group, energy in potential_energy1.items()}
        potential_energy2 = {group_to_force[group]: energy for group, energy in potential_energy2.items()}
        return potential_energy1, potential_energy2

    # Catch the exception and log table force type -> energy.
    try:
        potential_energy1, potential_energy2 = compare_context_energies(
            context1,
            context2,
            positions,
            box_vectors,
            by_force_group=by_force_type,
            rtol=rtol,
            atol=atol,
        )
    except FailedEnergyComparisonError as e:
        # We don't need to convert force groups into force types
        if not by_force_type:
            raise
        # Add to the error message the table of energies by force type.
        potential_energy1, potential_energy2 = map_energies_by_force_type(e.potential_energy1, e.potential_energy2)
        # Pretty-print the dictionaries.
        table = "\n\npotential energy system1:\n" + pprint.pformat(potential_energy1)
        table += "\n\npotential energy system2:\n" + pprint.pformat(potential_energy2)
        raise type(e)(str(e) + table, e.potential_energy1, e.potential_energy2)

    if by_force_type:
        return map_energies_by_force_type(potential_energy1, potential_energy2)
    return potential_energy1, potential_energy2


class _ParametersComparer:
    """This is just a convenience class to compare and print parameters.

    Parameters
    ----------
    **kwargs
        All the parameters to compare.

    Attributes
    ----------
    parameters : dict
        The dictionary of parameters.

    Examples
    --------
    >>> import copy
    >>> from openmm import unit
    >>> par1 = _ParametersComparer(charge=1.0*unit.elementary_charge,
    ...                            rmin_half=1.5*unit.angstrom)
    >>> par2 = copy.deepcopy(par1)
    >>> par3 = _ParametersComparer(charge=-1.0*unit.elementary_charge,
    ...                            rmin_half=0.15*unit.nanometers)
    >>> par1 == par2
    True
    >>> par1 == par3
    False
    >>> str(par1)
    '(charge: 1.0 e, rmin_half: 1.5 A)'

    """

    def __init__(self, **parameters):
        self.parameters = parameters

    def remove_parameters(self, parameter_names):
        for parameter_name in parameter_names:
            try:
                del self.parameters[parameter_name]
            except KeyError:
                pass

    def pretty_format_diff(self, other, new_line=True, indent=True):
        """Return a pretty-formatted string describing the differences between parameters.

        Parameters
        ----------
        other : _ParametersComparer
        new_line : bool, optional
            Separate different parameters with new lines.
        indent : bool, optional
            Indent the formatting.

        Returns
        -------
        diff_str : str

        """
        diff = self._get_diff(other)
        diff_list = []
        for par_name in sorted(diff.keys()):
            par1, par2 = diff[par_name]
            diff_list.append(f"{par_name}: {par1} != {par2}")
        separator = "\n" if new_line else ""
        diff_str = separator.join(diff_list)
        if indent:
            diff_str = textwrap.indent(diff_str, prefix="    ")
        return diff_str

    def _get_diff(self, other):
        """Build a 'diff' including only the parameters that are different."""
        assert set(self.parameters.keys()) == set(other.parameters.keys())
        diff = {}
        for par_name, par1_value in self.parameters.items():
            par2_value = other.parameters[par_name]
            # Determine whether this parameter is close or not.
            if not quantities_allclose(par1_value, par2_value):
                diff[par_name] = (par1_value, par2_value)
        return diff

    def __eq__(self, other):
        return len(self._get_diff(other)) == 0

    def __str__(self):
        # Reorder the parameters to print in deterministic order.
        parameter_order = sorted(self.parameters)
        # Quantities in a dictionary are normally printed in the
        # format "Quantity(value, unit=unit)" so we make it prettier.
        par_str = ", ".join(f"{p}: {self.parameters[p]}" for p in parameter_order)
        return "(" + par_str + ")"


class _TorsionParametersComparer:
    """A ParametersComparer class for torsions.

    Torsions require a different comparison because multiple set of
    (periodicity, phase, k) parameters can be associated to the same
    four atoms.

    Parameters
    ----------
    *parameters
        Variable length arguments. Each is a _ParametersComparer instance.

    Attributes
    ----------
    parameters : list[_ParametersComparer]
        The list of _ParametersComparer instances.

    Examples
    --------
    >>> from openmm import unit
    >>> par1 = _ParametersComparer(periodicity=2, phase=0.0*unit.degrees)
    >>> par2 = _ParametersComparer(periodicity=4, phase=180.0*unit.degrees)
    >>> torsion_par = _TorsionParametersComparer(par1, par2)

    """

    def __init__(self, *parameters):
        # *args is a tuple. Convert it to a list to make it appendable.
        self.parameters = list(parameters)

    def pretty_format_diff(self, other, new_line=True, indent=True):
        """Return a pretty-formatted string describing the differences between parameters.

        Parameters
        ----------
        other : _TorsionParametersComparer
        new_line : bool, optional
            Separate different parameters with new lines.
        indent : bool, optional
            Indent the formatting.

        Returns
        -------
        diff_str : str

        """
        diff_str = "Parameters in first system:\n"
        diff_str += self._pretty_format_parameters(self.parameters, new_line=new_line)
        diff_str += "\nParameters in second system:\n"
        diff_str += self._pretty_format_parameters(other.parameters, new_line=new_line)
        if indent:
            diff_str = textwrap.indent(diff_str, "    ")
        return diff_str

    def _pretty_format_parameters(self, parameters, new_line=True, indent=True):
        # Reorder the parameters by periodicity and then phase to print in deterministic order.
        def sort_key(x):
            return x.parameters["periodicity"], x.parameters["phase"]

        parameters = sorted(parameters, key=sort_key)
        # Quantities in a dictionary are normally printed in the
        # format "Quantity(value, unit=unit)" so we make it prettier.
        separator = "\n" if new_line else ", "
        prefix = "["
        if indent and new_line:
            separator += "    "
        elif indent:
            prefix += "    "
        return prefix + separator.join(str(pars) for pars in parameters) + "]"

    def __eq__(self, other):
        # Look for self.parameters that don't match any other.parameters.
        for parameters1 in self.parameters:
            # Find at least 1 set of parameters in the list that are equal.
            if all(parameters1 != parameters2 for parameters2 in other.parameters):
                return False

        # Look for other.parameters that don't match any other self.parameters.
        for parameters2 in other.parameters:
            # Find at least 1 set of parameters in the list that are equal.
            if all(parameters1 != parameters2 for parameters1 in self.parameters):
                return False
        return True

    def __str__(self):
        return self._pretty_format_parameters(self.parameters, new_line=False, indent=False)


class FailedParameterComparisonError(AssertionError):
    """Error raised when the parameter comparison between two forces fails.

    Attributes
    ----------
    different_parameters : dict[Hashable, tuple[_ParameterComparer]]
        All the parameters for which a differences was detected.
        different_parameters[atom_indices] is a tuple
        (parameter_force1, parameter_force2).

    """

    def __init__(self, err_msg, different_parameters):
        super().__init__(err_msg)
        self.different_parameters = different_parameters


def _compare_parameters(
    parameters_force1,
    parameters_force2,
    interaction_type,
    force_name="",
    systems_labels=None,
):
    """Compare the parameters of 2 forces and raises an exception if they are different.

    Parameters
    ----------
    parameters_force1 : dict[Hashable, _ParametersComparer]
        A dictionary associating a _ParametersComparer entry to a key
        that is usually a set of atom indices (e.g., the atom index
        for nonbonded interactions, two atom indices for bonds and
        1-4 exceptions, 3 indices for angles, and 4 for torsions).
    parameters_force2 : dict[Hashable, _ParametersComparer]
        The parameters of the force to compare.
    interaction_type : str
        A string describing the type of interactions (e.g., "bond",
        "particle exception", "proper torsion"). This is only used to
        improve the error message where differences between the two
        forces are detected.
    force_name : str, optional
        The name of the force to optionally include in the eventual
        error message.
    systems_labels : tuple[str], optional
        A pair of strings with a meaningful name for the system. If
        specified, this will be included in the error message to
        improve its readability.

    Raises
    ------
    FailedParameterComparisonError
        If there are differences in the parameters of the two forces.
        The exceptions exposes an attribute ``different_parameters``
        with the different parameters.

    """
    diff_msg = ""

    # First check the parameters that are unique to only one of the forces.
    unique_keys1 = set(parameters_force1) - set(parameters_force2)
    unique_keys2 = set(parameters_force2) - set(parameters_force1)
    err_msg = "\n\n{} force has the following unique " + interaction_type + "s: {}\n"
    if len(unique_keys1) != 0:
        force_label = systems_labels[0] if systems_labels is not None else "First"
        diff_msg += err_msg.format(force_label, sorted(unique_keys1))
    if len(unique_keys2) != 0:
        force_label = systems_labels[1] if systems_labels is not None else "Second"
        diff_msg += err_msg.format(force_label, sorted(unique_keys2))

    # Create a diff for entries that have same keys but different parameters.
    different_parameters = {}
    for key in set(parameters_force1).intersection(set(parameters_force2)):
        if parameters_force1[key] != parameters_force2[key]:
            different_parameters[key] = (parameters_force1[key], parameters_force2[key])

    # Handle force and systems labels default arguments.
    if force_name != "":
        force_name += " "  # Add space after.
    if systems_labels is not None:
        systems_labels = " for the {} and {} systems respectively".format(*systems_labels)
    else:
        systems_labels = ""

    # Print error.
    if len(different_parameters) > 0:
        diff_msg += (
            f"\n\nThe following {interaction_type}s have different parameters "
            f"in the two {force_name}forces{systems_labels}:"
        )
        for key in sorted(different_parameters.keys()):
            param1, param2 = different_parameters[key]
            parameters_diff = param1.pretty_format_diff(param2)
            diff_msg += f"\n{interaction_type} {key}:\n{parameters_diff}"

    if diff_msg != "":
        diff_msg = (f"A difference between {interaction_type} was detected. Details follow.") + diff_msg
        raise FailedParameterComparisonError(diff_msg + "\n", different_parameters)


@functools.singledispatch
def _get_force_parameters(force, system, ignored_parameters):
    """This function builds a _ParameterComparer representation of the
    force parameters that can be used for comparison to other forces
    with _compare_parameters.

    Each _ParameterComparer must be associated to a key, which is usually
    a tuple of atom indices (e.g., the atom index for nonbonded interactions,
    two atom indices for bonds and 1-4 exceptions, 3 indices for angles,
    and 4 for torsions).

    A single force can generate multiple representations (e.g., a NonbondedForce
    can generate one representations for particle parameters and one for
    parameter exceptions.

    Parameters
    ----------
    force
        The OpenMM Force object.
    system : openmm.System
        The System to which this force belongs to.
    ignored_parameters : Iterable[str]
        The parameters to ignore in the comparison.

    Returns
    -------
    force_parameters : dict[str, dict[Hashable, _ParameterComparer]]
        The parameter representation. force_parameters[interaction_type]
        is a dictionary mapping a tuple of atom indices to the
        associated parameters. interaction_type can be any string
        identifying the type of parameters (e.g. "particle exception",
        "proper torsion", "bond").

    """
    raise NotImplementedError(f"Comparison between {type(force)}s is not currently supported.")


@_get_force_parameters.register(openmm.HarmonicBondForce)
def _get_bond_force_parameters(force, _, ignored_parameters):
    """Implementation of _get_force_parameters for HarmonicBondForces."""
    # Build the dictionary of the parameters of a single force.
    force_parameters = {}
    for bond_idx in range(force.getNumBonds()):
        atom1, atom2, r0, k = force.getBondParameters(bond_idx)

        # Ignore bonds with 0.0 spring constant that don't affect the energy.
        if k._value == 0.0:
            continue

        # Reorder the bond to have a canonical key.
        bond_key = tuple(sorted([atom1, atom2]))
        force_parameters[bond_key] = _ParametersComparer(r0=r0, k=k)
        force_parameters[bond_key].remove_parameters(ignored_parameters)

    return {"bond": force_parameters}


@_get_force_parameters.register(openmm.HarmonicAngleForce)
def _get_angle_force_parameters(force, _, ignored_parameters):
    """Implementation of _get_force_parameters for HarmonicAngleForces."""
    # Build the dictionary of the parameters of a single force.
    force_parameters = {}
    for angle_idx in range(force.getNumAngles()):
        atom1, atom2, atom3, theta0, k = force.getAngleParameters(angle_idx)

        # Ignore angles with 0.0 spring constant that don't affect the energy.
        if k._value == 0.00:
            continue

        # Reorder the bond to have a canonical key.
        angle_key = min(atom1, atom3), atom2, max(atom1, atom3)
        force_parameters[angle_key] = _ParametersComparer(theta0=theta0, k=k)
        force_parameters[angle_key].remove_parameters(ignored_parameters)

    return {"angle": force_parameters}


@_get_force_parameters.register(openmm.NonbondedForce)
def _get_nonbonded_force_parameters(force, _, ignored_parameters):
    """Implementation of _get_force_parameters for NonbondedForces."""
    # Build the dictionary of the particle parameters of the force.
    particle_parameters = {}
    for particle_idx in range(force.getNumParticles()):
        charge, sigma, epsilon = force.getParticleParameters(particle_idx)

        # Ignore sigma parameter if epsilon is 0.0.
        particle_parameters[particle_idx] = _ParametersComparer(charge=charge, epsilon=epsilon)
        if epsilon._value != 0.0:
            particle_parameters[particle_idx].parameters["sigma"] = sigma
        particle_parameters[particle_idx].remove_parameters(ignored_parameters)

    # Build the dictionary representation of the particle exceptions.
    exception_parameters = {}
    for exception_idx in range(force.getNumExceptions()):
        atom1, atom2, chargeprod, sigma, epsilon = force.getExceptionParameters(exception_idx)

        # Reorder the atom indices to have a canonical key.
        exception_key = tuple(sorted([atom1, atom2]))
        # Ignore sigma parameter if epsilon is 0.0.
        exception_parameters[exception_key] = _ParametersComparer(charge=chargeprod, epsilon=epsilon)
        if epsilon._value != 0.0:
            exception_parameters[exception_key].parameters["sigma"] = sigma
        exception_parameters[exception_key].remove_parameters(ignored_parameters)

    return {"particle": particle_parameters, "particle exception": exception_parameters}


def _merge_impropers_folds(improper_parameters):
    """Merge impropers with same periodicities and phases into a single folds.

    See _get_torsion_force_parameters.

    """
    for improper_comparer in improper_parameters.values():
        # Group parameters by periodicity and phase and sum force constants.
        for comparer2_idx in reversed(range(1, len(improper_comparer.parameters))):
            parameter_comparer2 = improper_comparer.parameters[comparer2_idx]
            parameters2 = parameter_comparer2.parameters

            # parameter_comparer1 comes before parameter_comparer2 in the list.
            for parameter_comparer1 in improper_comparer.parameters[:comparer2_idx]:
                parameters1 = parameter_comparer1.parameters

                # If phase and periodicity match, sum the force constants
                # and remove one _ParametersComparer object from the list.
                if (parameters1["phase"] == parameters2["phase"]) and (
                    parameters1["periodicity"] == parameters2["periodicity"]
                ):
                    # Update the force constant of the first parameter.
                    parameters1["k"] += parameters2["k"]
                    # Remove the second parameter from the list.
                    del improper_comparer.parameters[comparer2_idx]
                    break


@_get_force_parameters.register(openmm.PeriodicTorsionForce)
def _get_torsion_force_parameters(force, system, ignored_parameters):
    """Implementation of _get_force_parameters for NonbondedForces.

    If the special attribute 'n_improper_folds' is in ignored_parameters,
    then impropers with same periodicity and phases are merged into a
    single set of parameters having the sum of the force constants as k.
    This allows to compare parameters before and after the 0.1 version of
    the toolkit, in which the six-fold implementation was changed into
    the current three-fold.

    """
    if "n_improper_folds" not in ignored_parameters:
        ignore_n_folds = False
    else:
        ignore_n_folds = True
        # Create a copy of ignored_parameters without n_improper_folds
        # that must be removed later from the _ParametersComparer object.
        ignored_parameters = copy.deepcopy(ignored_parameters)
        ignored_parameters.remove("n_improper_folds")

    # Find all bonds. We'll use this to distinguish
    # between proper and improper torsions.
    bond_set = _find_all_bonds(system)

    proper_parameters = {}
    improper_parameters = {}
    for torsion_idx in range(force.getNumTorsions()):
        atom1, atom2, atom3, atom4, periodicity, phase, k = force.getTorsionParameters(torsion_idx)

        # Ignore torsions that don't contribute to the energy.
        if k._value == 0.0:
            continue

        torsion_key = [atom1, atom2, atom3, atom4]
        if len(set(torsion_key)) != 4:
            raise ValueError(f"Torsion {torsion_idx} is defined on less than 4 atoms: {torsion_key}")

        # Check if this is proper or not.
        torsion_bonds = [(torsion_key[i], torsion_key[i + 1]) for i in range(3)]
        is_proper = all(bond in bond_set for bond in torsion_bonds)

        # Determine the canonical order of the torsion key.
        if is_proper:
            torsion_key = _get_proper_torsion_canonical_order(*torsion_key)
            force_parameters = proper_parameters
        else:
            torsion_key = _get_improper_torsion_canonical_order(bond_set, *torsion_key)
            force_parameters = improper_parameters

        # Update the dictionary representation.
        parameters = _ParametersComparer(periodicity=periodicity, phase=phase, k=k)
        parameters.remove_parameters(ignored_parameters)
        # Each set of 4 atoms can have multiple torsion terms.
        try:
            force_parameters[torsion_key].parameters.append(parameters)
        except KeyError:
            force_parameters[torsion_key] = _TorsionParametersComparer(parameters)

    # Sum the force constants of all the improper torsion folds. This
    # is fundamental to compare parameters with versions of the toolkit
    # previous to 0.1, where the trefoil improper was six-fold rather
    # than three-fold as in the current version.
    if ignore_n_folds:
        _merge_impropers_folds(improper_parameters)

    return {
        "proper torsion": proper_parameters,
        "improper torsion": improper_parameters,
    }


def _find_all_bonds(system):
    """Find all bonds in the given system

    Parameters
    ----------
    system : openmm.System
        A System object containing a HarmonicBondForce from which the
        bonds are inferred.

    Returns
    -------
    bond_set : set[tuple[int]]
        A set of pairs (atom_index1, atom_index2) associated to bonds.
        For each bond two pairs are created with the atom indices in
        inverse order.

    """
    # Find the force with information on the bonds.
    bond_force = [f for f in system.getForces() if isinstance(f, openmm.HarmonicBondForce)]
    assert len(bond_force) == 1
    bond_force = bond_force[0]

    # Create a set of all bonds.
    bond_set = set()
    for bond_idx in range(bond_force.getNumBonds()):
        atom1, atom2, _, _ = bond_force.getBondParameters(bond_idx)
        bond_set.add((atom1, atom2))
        bond_set.add((atom2, atom1))

    # Find all the constrained bonds, which do not have a HarmonicBondForce term.
    for constraint_idx in range(system.getNumConstraints()):
        atom1, atom2, _ = system.getConstraintParameters(constraint_idx)
        bond_set.add((atom1, atom2))
        bond_set.add((atom2, atom1))

    return bond_set


def _get_proper_torsion_canonical_order(i0, i1, i2, i3):
    """Create a unique order of the 4 atom indices of a proper torsion.

    The atom indices of a proper torsion are reordered so that the
    first atom is the smallest index.

    Parameters
    ----------
    i0, i1, i2, i3 : int
        Atom indices of the proper torsion.

    Returns
    -------
    j0, j1, j2, j3 : int
        Reordered atom indices of the proper torsion.

    """
    if i0 < i3:
        return i0, i1, i2, i3
    else:
        return i3, i2, i1, i0


def _get_improper_torsion_canonical_order(bond_set, i0, i1, i2, i3):
    """Create a unique order of the 4 atom indices of an improper torsion.

    Return j0, j1, j2, j3, where j0 is the central index and j1, j2, j3
    are in sorted() order. The centrality is determined by the maximum
    counts in the adjacency matrix.

    Parameters
    ----------
    bond_set : set
        The set of all bonds as determined by _find_all_bonds.
    i0, i1, i2, i3 : int,
        Atom indices of the improper torsion.

    Returns
    -------
    j0, j1, j2, j3 : int
        Reordered atom indices of the improper torsion, with j0 being
        the central index.

    """
    connections = np.zeros((4, 4))

    mapping = {i0: 0, i1: 1, i2: 2, i3: 3}
    inv_mapping = dict([(val, key) for key, val in mapping.items()])

    for a, b in itertools.combinations([i0, i1, i2, i3], 2):
        if (a, b) in bond_set:
            i, j = mapping[a], mapping[b]
            connections[i, j] += 1.0
            connections[j, i] += 1.0

    central_ind = connections.sum(0).argmax()
    central_ind = inv_mapping[central_ind]
    other_ind = sorted([i0, i1, i2, i3])
    other_ind.remove(central_ind)

    return central_ind, other_ind[0], other_ind[1], other_ind[2]


def compare_system_parameters(
    system1,
    system2,
    systems_labels=None,
    ignore_charges=False,
    ignore_improper_folds=False,
):
    """Check that two OpenMM systems have the same parameters.

    Parameters
    ----------
    system1 : openmm.System
        The first system to compare.
    system2 : openmm.System
        The second system to compare.
    systems_labels : tuple[str], optional
        A pair of strings with a meaningful name for the system. If
        specified, this will be included in the error message to
        improve its readability.
    ignore_charges : bool, optional
        If True, particle and exception charges are ignored during
        the comparison. Default is False.
    ignore_improper_folds : bool, optional
        If True, all the folds of the same impropers are merged into
        a single one for the purpose of parameter comparison. Default
        is False.

    Raises
    ------
    FailedParameterComparisonError
        If there are differences in the parameters of the two forces.
        The exceptions exposes an attribute ``different_parameters``
        with the different parameters.

    """
    # Determine parameters to ignore.
    ignored_parameters_by_force = {}
    if ignore_charges:
        ignored_parameters_by_force["NonbondedForce"] = ["charge", "chargeprod"]
    if ignore_improper_folds:
        ignored_parameters_by_force["PeriodicTorsionForce"] = ["n_improper_folds"]

    # We need to perform some checks on the type and number of forces in the Systems.
    force_names1 = collections.Counter(f.__class__.__name__ for f in system1.getForces())
    force_names2 = collections.Counter(f.__class__.__name__ for f in system2.getForces())
    err_msg = "Only systems having 1 force per type are supported."
    assert set(force_names1.values()) == {1}, err_msg
    assert set(force_names2.values()) == {1}, err_msg

    # Remove the center-of-mass motion remover force from comparison.
    for force_names in [force_names1, force_names2]:
        if "CMMotionRemover" in force_names:
            del force_names["CMMotionRemover"]

    # Check that the two systems have the same forces.
    if set(force_names1) != set(force_names2):
        err_msg = "The two Systems have different forces to compare:\n - system1 {}\n - system2 {}"
        raise ValueError(err_msg.format(sorted(force_names1), sorted(force_names2)))

    # Find all the pair of forces to compare.
    force_pairs = {force_name: [] for force_name in force_names1}
    for system in [system1, system2]:
        for force in system.getForces():
            # Skip CMMotionRemover.
            if force.__class__.__name__ == "CMMotionRemover":
                continue
            force_pairs[force.__class__.__name__].append(force)

    # TODO: ParmEd-derived system representations separate out vdW from ES (putting ES in NonbondedForce
    # and vdW in CustomNonbondedForce with precomputed mixing parameters). We'll need to update this to
    # somehow account for the difference.

    # Compare all pairs of forces
    for force_name, (force1, force2) in force_pairs.items():
        # Select the parameters to ignore.
        try:
            ignored_parameters = ignored_parameters_by_force[force_name]
        except KeyError:
            ignored_parameters = ()

        # Get a dictionary representation of the parameters.
        parameters_force1 = _get_force_parameters(force1, system1, ignored_parameters)
        parameters_force2 = _get_force_parameters(force2, system2, ignored_parameters)

        # Compare the forces and raise an error if a difference is detected.
        for parameter_type, parameters1 in parameters_force1.items():
            parameters2 = parameters_force2[parameter_type]
            _compare_parameters(
                parameters1,
                parameters2,
                interaction_type=parameter_type,
                force_name=force_name,
                systems_labels=systems_labels,
            )


def compare_amber_smirnoff(
    prmtop_file_path,
    inpcrd_file_path,
    forcefield,
    molecule,
    check_parameters=True,
    check_energies=True,
    **kwargs,
):
    """
    Compare energies and parameters for OpenMM Systems/topologies created
    from an AMBER prmtop and crd versus from a SMIRNOFF force field file which
    should parameterize the same system with same parameters.

    Parameters
    ----------
    prmtop_file_path : str
        Path to the topology/parameter file in AMBER prmtop format
    inpcrd_file_path : str
        Path to the coordinates file in AMBER inpcrd or rst7 format
    forcefield : ForceField
        Force field instance used to create the system to compare.
    molecule : topology.molecule.Molecule
        The molecule object to test.
    check_parameters : bool, optional
        If False, parameters are not compared. Energies are still comapred
        if ``check_energies`` is ``True`` (default is ``True``).
    check_parameters : bool, optional
        If False, energies are not compared and they are not returned.
        Parameters are not compared if ``check_energies`` is ``True``
        (default is ``True``).
    ignore_charges : bool, optional
        If True, particle and exception charges are ignored during
        the comparison. Default is False.

    Returns
    -------
    energies : dict[str, dict[str, openmm.unit.Quantity]] or None
        If ``check_energies`` is ``False``, nothing is returned. Otherwise,
        this is a dictionary with two keys: 'AMBER' and 'SMIRNOFF', each
        pointing to another dictionary mapping forces to their total
        contribution to the potential energy.

    Raises
    ------
    FailedEnergyComparisonError
        If the potential energies of the two systems are different
        or if the two Systems in the contexts have forces divided into
        different force groups. The potential energies of the AMBER and
        ForceField systems can be accessed in the Exception through the
        attributes potential_energy1 and potential_energy2 respectively.
    FailedParameterComparisonError
        If there are differences in the parameters of the two systems.
        The exceptions exposes an attribute ``different_parameters``
        with the different parameters.

    """
    from openff.toolkit.topology import Topology

    # Create System from AMBER files. By default, contrarily to ForceField,
    # systems from AMBER files are created with removeCMMotion=True
    amber_system, openmm_topology, positions = create_system_from_amber(
        prmtop_file_path, inpcrd_file_path, removeCMMotion=False
    )
    box_vectors = amber_system.getDefaultPeriodicBoxVectors()

    # Create System from force field.
    openff_topology = Topology.from_openmm(openmm_topology, unique_molecules=[molecule])
    ff_system = forcefield.create_openmm_system(openff_topology)

    # Test energies and parameters.
    if check_parameters:
        compare_system_parameters(amber_system, ff_system, systems_labels=("AMBER", "SMIRNOFF"), **kwargs)

    if check_energies:
        amber_energies, forcefield_energies = compare_system_energies(
            amber_system, ff_system, positions, box_vectors, **kwargs
        )

        return {"AMBER": amber_energies, "SMIRNOFF": forcefield_energies}
    return None


def reorder_openff_to_openmm(xyz, n_atoms_per_mol, n_vptls_per_mol):
    """
    Rearrange the coordinates of an OpenFF-created system to an order
    that OpenMM uses/used.

    This function is mainly used to compare with a system that was created with
    OpenMM, where the ordering is AAABBAAABB and was created elsewhere. This will
    reorder the OpenFF particles such that they match with this ordering.

    xyz : list[list[float]]
        The atom and virtual positions
    n_atoms_per_mol : int
        The number of atoms per molecule. All molecules are assumed to have the
        same number of atoms.
    n_vptls_per_mol : int
        The number of virtual particles per molecule. All molecules are assumed
        to have the same number of virtual particles

    Returns:
    xyz : list[list[float]]
        The atom and virtual particle positions in OpenFF order
    """

    n_particles_per_mol = n_atoms_per_mol + n_vptls_per_mol
    n_mols = xyz.shape[0] // n_particles_per_mol

    atom_base_indices = np.arange(n_atoms_per_mol)
    vsite_base_indices = np.arange(n_vptls_per_mol) + n_atoms_per_mol * n_mols

    mask = np.hstack(
        [
            np.hstack(
                [
                    atom_base_indices + (i * n_atoms_per_mol),
                    vsite_base_indices + (i * n_vptls_per_mol),
                ]
            )
            for i in range(n_mols)
        ]
    )

    return xyz[mask]


def reorder_openmm_to_openff(xyz, n_atoms_per_mol, n_vptls_per_mol):
    """
    Rearrange the coordinates of an OpenMM-created system to an order
    that OpenFF uses.

    For example, since OpenMM creates virtual sites in-place, they are interleaved
    with the atoms. This means that, if atoms are labeled as A and virtual sites
    as B, then the ordering will be AAABBAAABB, and the index is 0 to 9.
    This function will reorder the coordinates to be AAAAAABBBB, which is the order
    OpenFF expects/uses.

    xyz : list[list[float]]
        The atom and virtual positions
    n_atoms_per_mol : int
        The number of atoms per molecule. All molecules are assumed to have the
        same number of atoms.
    n_vptls_per_mol : int
        The number of virtual particles per molecule. All molecules are assumed
        to have the same number of virtual particles

    Returns:
    xyz : list[list[float]]
        The atom and virtual particle positions in OpenFF order
    """

    # constants for simplicity
    n_particles_per_mol = n_atoms_per_mol + n_vptls_per_mol
    n_mols = xyz.shape[0] // n_particles_per_mol

    atom_base_indices = np.arange(n_atoms_per_mol)
    atom_indices = np.hstack([atom_base_indices + (i * n_particles_per_mol) for i in range(n_mols)])

    vsite_base_indices = np.arange(n_vptls_per_mol) + n_atoms_per_mol
    vsite_indices = np.hstack([vsite_base_indices + (i * n_particles_per_mol) for i in range(n_mols)])

    mask = np.hstack([atom_indices, vsite_indices])

    return xyz[mask]


def openmm_evaluate_vsites_and_energy(omm_top, omm_sys, atom_xyz_in_nm, minimize=False):
    """
    Calculate the Virtual Site positions and potential energy
    of the given system

    Parameters
    ----------
    omm_top: OpenMM Topology
    omm_sys: OpenMM System
    atom_xyz_in_nm: list[list[float]]
        An N,3 array of floats interpreted as nanometers of the atom positions
    minimize: bool
        Perform a minimization before calculating energy

    Returns
    -------
    pos: list[list[float]]
        The coordinates of all particles in the system (OpenMM ordering)
    ene: float
        The potential energy
    """

    integ = openmm.VerletIntegrator(1.0 * openmm_unit.femtoseconds)

    sim = openmm.app.Simulation(omm_top, omm_sys, integ)
    pos = sim.context.getState(getPositions=True).getPositions()

    sim.context.setPositions(atom_xyz_in_nm)

    # Once the atom positions are known, virtual sites can
    # be determined.
    sim.context.computeVirtualSites()

    if minimize:
        sim.minimizeEnergy()

    state = sim.context.getState(getEnergy=True, getPositions=True)
    ene = state.getPotentialEnergy()
    pos = [list(xyz) for xyz in state.getPositions().value_in_unit(openmm_unit.nanometer)] * unit.nanometer
    return pos, ene


def insert_vsite_padding(xyz, n_mol, n_vptl, order="OpenFF"):
    """
    Given a set of atomic coordinates, insert dummy coordinates in indices
    expected for either OpenFF or OpenMM. OpenFF places the dummy coordinates
    at the end, whereas OpenMM expects virtual sites to be next to their owning
    molecule.

    Parameters
    ----------
    xyz: list[list[float]]
        An Nx3 set of coordinates of atoms only
    n_mol: int
        The number of molecules
    n_vptl: int
        The number of virtual site particle dummy positions to insert

    Returns
    -------
    xyz : list[list[float]]
        The coordinates of shape (N+n_vsite,3) with all particles in the expected order
    """

    # heuristic that vsites come at the end; applies to OpenFF
    vsite_pad = np.arange(n_vptl * 3).reshape(n_vptl, 3)
    xyz = np.vstack([xyz, vsite_pad])

    # but does not apply to OpenMM. Need to interleave vsite coordinates
    if order == "OpenMM":
        n_vptls_per_mol = n_vptl // n_mol
        n_atoms_per_mol = len(xyz) // n_mol - n_vptls_per_mol
        xyz = reorder_openff_to_openmm(xyz, n_atoms_per_mol, n_vptls_per_mol)

    return xyz


def coords_from_off_mols(mols, conformer_id=0, unit=unit.angstrom):
    """
    Small utility to extract the coordinates from the molecule conformers,
    with optional unit conversion

    Parameters
    ----------
    mols : list[openff.toolkit.topology.molecule.Molecule]
        The system/list of molecules
    conformer_id : int, default=0
        The conformer to use for retreiving the coordinates from each molcule
    unit : openmm.unit, default=unit.angstrom
        The units to convert the coordinates to
    """
    xyz = np.vstack([mol.conformers[conformer_id].m_as(unit) for mol in mols])
    return xyz * unit


def evaluate_molecules_omm(water, ff, minimize=False):
    """
    Given a list of molecules and a force field definition, calculate the
    positions and energy.

    Parameters
    ----------
    molecules: list[openff.toolkit.topology.molecule.Molecule]
        A list of molecules with a 3D conformation
    forcefield: openff.toolkit.typing.engines.smirnoff.forcefield.ForceField
        The force field object to parameterize with
    minimize: bool
        Whether the structure should be minimized

    Returns
    -------
    xyz: list[list[float]]
        The coordinates of all particles in the system (OpenMM ordering)
    ene: float
        The potential energy
    """

    from openff.toolkit.topology import Topology

    top = Topology.from_molecules(water).to_openmm()

    # get the oFF topology and the positions
    atom_xyz = to_openmm(coords_from_off_mols(water, unit=unit.nanometer))

    # first pass; no virtual sites
    omm_modeller = openmm.app.Modeller(top, atom_xyz)
    omm_modeller.addExtraParticles(ff)

    # second pass: virtual sites now present (according to oMM FF)
    top = omm_modeller.getTopology()

    sys = ff.createSystem(top, nonbondedMethod=openmm.app.NoCutoff)
    n_vptl = sys.getNumParticles() - len(atom_xyz)
    n_mol = len(water)

    # need to insert dummy positions into coordinates, since OpenMM
    # systems create virtual site particles in place, i.e. OHHMMOHHMM etc.
    atom_xyz = insert_vsite_padding(atom_xyz, n_mol, n_vptl, order="OpenMM")

    # Need an openMM topology, an openMM system, and a set of positions
    # to calculate the energy. Also returns the vsite positions based
    # on supplied atom coordinates
    xyz, ene = openmm_evaluate_vsites_and_energy(top, sys, atom_xyz, minimize=minimize)

    return xyz, ene


def evaluate_molecules_off(molecules, forcefield, minimize=False):
    """
    Given a list of molecules and a force field definition, calculate the
    positions and energy.

    Parameters
    ----------
    molecules: list[openff.toolkit.topology.molecule.Molecule]
        A list of molecules with a 3D conformation
    forcefield: openff.toolkit.typing.engines.smirnoff.forcefield.ForceField
        The force field object to parameterize with
    minimize: bool
        Whether the structure should be minimized

    Returns
    -------
    xyz: list[list[float]]
        The coordinates of all particles in the system (OpenMM ordering)
    ene: float
        The potential energy
    """

    from openff.toolkit.topology import Topology

    # get the oFF topology and the positions
    top = Topology.from_molecules(molecules)
    atom_xyz = to_openmm(coords_from_off_mols(molecules, unit=unit.nanometer))

    # IMPORTANT! The new system has the virtual sites, but the oFF top does not as there is no API yet to do so
    sys = forcefield.create_openmm_system(top, allow_nonintegral_charges=True)

    n_vptl = sys.getNumParticles() - len(atom_xyz)
    n_mol = len(molecules)
    atom_xyz = insert_vsite_padding(atom_xyz, n_mol, n_vptl, order="OpenFF")

    # Need an openMM topology, an openMM system, and a set of positions
    # to calculate the energy. Also returns the vsite positions based
    # on supplied atom coordinates
    xyz, ene = openmm_evaluate_vsites_and_energy(top.to_openmm(), sys, atom_xyz, minimize=minimize)

    return xyz, ene


def get_14_scaling_factors(omm_sys: openmm.System) -> tuple[list, list]:
    """Find the 1-4 scaling factors as they are applied to an OpenMM System"""
    nonbond_force = next(f for f in omm_sys.getForces() if type(f) is openmm.NonbondedForce)

    vdw_14 = list()
    coul_14 = list()

    for exception_idx in range(nonbond_force.getNumExceptions()):
        i, j, q, _, eps = nonbond_force.getExceptionParameters(exception_idx)

        # Trust that q == 0 covers the cases of 1-2, 1-3, and truly being 0
        if q._value != 0:
            q_i = nonbond_force.getParticleParameters(i)[0]
            q_j = nonbond_force.getParticleParameters(j)[0]
            coul_14.append(q / (q_i * q_j))

        if eps._value != 0:
            eps_i = nonbond_force.getParticleParameters(i)[2]
            eps_j = nonbond_force.getParticleParameters(j)[2]
            vdw_14.append(eps / (eps_i * eps_j) ** 0.5)

    return coul_14, vdw_14


def compare_partial_charges(system1, system2):
    assert system1.getNumParticles() == system2.getNumParticles(), "Systems do not have the same number of particles"

    n_particles = system1.getNumParticles()

    for force1 in system1.getForces():
        if type(force1) is openmm.NonbondedForce:
            charges1 = [force1.getParticleParameters(i)[0] for i in range(n_particles)]

    for force2 in system2.getForces():
        if type(force2) is openmm.NonbondedForce:
            charges2 = [force2.getParticleParameters(i)[0] for i in range(n_particles)]

    assert charges1 == charges2, [(x - y) for x, y in zip(charges1, charges2)]
