"""Tests reproducing specific issues that are otherwise uncategorized."""

import random
import warnings

import numpy
import pytest
from openff.packmol import UNIT_CUBE, pack_box, solvate_topology
from openff.toolkit import ForceField, Molecule, Quantity, Topology
from openff.utilities import get_data_file_path, has_executable, skip_if_missing

from openff.interchange import Interchange
from openff.interchange._tests import MoleculeWithConformer, _rng, needs_gmx, shuffle_topology
from openff.interchange.drivers import get_amber_energies, get_gromacs_energies, get_openmm_energies
from openff.interchange.exceptions import NonperiodicNoCutoffNotSupportedError
from openff.interchange.warnings import PresetChargesAndVirtualSitesWarning


def test_issue_723():
    parmed = pytest.importorskip("parmed")
    force_field = ForceField("openff-2.1.0.offxml")

    molecule = Molecule.from_smiles("C#N")
    molecule.generate_conformers(n_conformers=1)

    force_field.create_interchange(molecule.to_topology()).to_top("_x.top")

    parmed.load_file("_x.top")


@pytest.mark.skipif(not has_executable("packmol"), reason="Packmol is not installed")
@pytest.mark.parametrize("pack", [True, False])
def test_issue_1022(pack):
    topology = Topology.from_molecules(
        [
            MoleculeWithConformer.from_smiles(smi)
            for smi in [
                "CBr",
                "O",
                "O",
                "O",
                "[Na+]",
                "[Cl-]",
            ]
        ],
    )

    topology.box_vectors = Quantity(numpy.eye(3) * 10.0, "nanometer")

    if pack:
        topology = pack_box(
            molecules=[*topology.unique_molecules],
            number_of_copies=[1, 3, 1, 1],
            box_vectors=topology.box_vectors,
        )

    force_field = ForceField(
        "openff-2.0.0.offxml",
        get_data_file_path(
            "example-sigma-hole-bromine.offxml",
            "openff.interchange._tests.data",
        ),
    )

    interchange = force_field.create_interchange(topology)

    interchange.to_top("tmp")

    if pack:
        for seed in random.sample(range(0, 10**10), 5):
            # TODO: Compare GROMACS energies here as well
            get_openmm_energies(interchange).compare(
                get_openmm_energies(
                    shuffle_topology(
                        interchange,
                        force_field,
                        seed=seed,
                    ),
                ),
                tolerances={"Nonbonded": Quantity("1e-3 kilojoule_per_mole")},
            )


@skip_if_missing("openmm")
def test_issue_1031():
    import openmm.app

    # just grab some small PDB file from the toolkit, doesn't need to be huge, just
    # needs to include some relevant atom names
    openmm_topology = openmm.app.PDBFile(
        get_data_file_path(
            "proteins/MainChain_HIE.pdb",
            "openff.toolkit",
        ),
    ).topology

    openmm_atom_names = {atom.name for atom in openmm_topology.atoms()}

    interchange = Interchange.from_openmm(
        system=openmm.app.ForceField(
            "amber99sb.xml",
            "tip3p.xml",
        ).createSystem(
            openmm_topology,
            nonbondedMethod=openmm.app.PME,
        ),
        topology=openmm_topology,
    )

    openff_atom_names = {atom.name for atom in interchange.topology.atoms}

    assert sorted(openmm_atom_names) == sorted(openff_atom_names)

    # check a few atom names to ensure these didn't end up being empty sets
    for atom_name in ("NE2", "H3", "HA", "CH3", "CA", "CB", "CE1"):
        assert atom_name in openff_atom_names


def test_issue_1049():
    pytest.importorskip("openmm")

    topology = Topology.from_molecules(
        [
            Molecule.from_smiles("C"),
            Molecule.from_smiles("O"),
            Molecule.from_smiles("O"),
        ],
    )

    interchange = ForceField("openff-2.2.0.offxml", "opc.offxml").create_interchange(topology)

    openmm_topology = interchange.to_openmm_topology()
    openmm_system = interchange.to_openmm_system()

    # the same index in system should also be a virtual site in the topology
    for particle_index, particle in enumerate(openmm_topology.atoms()):
        assert openmm_system.isVirtualSite(particle_index) == (particle.element is None), (
            f"particle index {particle_index} is a virtual site in the system OR topology but not both"
        )


def test_issue_1052(sage, ethanol):
    """Test that _SMIRNOFFElectrostaticsCollection.charges is populated."""
    out = sage.create_interchange(ethanol.to_topology())

    assert len(out["Electrostatics"].charges) > 0


@skip_if_missing("openmm")
def test_issue_1209(sage, ethanol):
    ethanol.assign_partial_charges("gasteiger")
    ethanol.generate_conformers(n_conformers=1)

    out = sage.create_interchange(ethanol.to_topology(), charge_from_molecules=[ethanol])

    get_openmm_energies(out, combine_nonbonded_forces=False).compare(
        get_openmm_energies(
            Interchange.model_validate_json(out.model_dump_json()),
            combine_nonbonded_forces=False,
        ),
    )


def test_issue_1337(water):
    """Reproduce Issue #1337."""
    ff14sb_tip3p = ForceField("ff14sb_off_impropers_0.0.3.offxml", "tip3p.offxml")

    # just make sure this doesn't crash (original issue)
    # more extensive tests in openff/interchange/_tests/unit_tests/interop/openmm/test_constraints.py
    ff14sb_tip3p.create_interchange(water.to_topology())


@needs_gmx
def test_issue_1361_gromacs(caffeine, sage, tmp_path):
    """Test that the 'how to opt in to pseudo-vacuum' message is communicated."""
    interchange = sage.create_interchange(caffeine.to_topology())

    with pytest.raises(
        NonperiodicNoCutoffNotSupportedError,
        match=r"GROMACS versions 2020 and newer do not support systems without periodicity",
    ):
        interchange.to_gromacs(prefix="foo")


@pytest.mark.skipif(not has_executable("sander"), reason="sander not installed")
def test_issue_1361_amber(caffeine, sage, tmp_path):
    """Test that the 'how to opt in to pseudo-vacuum' message is communicated."""
    interchange = sage.create_interchange(caffeine.to_topology())

    with pytest.raises(
        NonperiodicNoCutoffNotSupportedError,
        match=r"vdW method no-cutoff not supported",
    ):
        interchange.to_amber(prefix="bar")


@pytest.mark.skipif(not has_executable("sander"), reason="sander not installed")
def test_issue_1395_amber(caffeine, sage, water, tmp_path):
    """Test that water charges are correctly ordered in ligand + water systems to Amber."""

    # turn off switching to make energy comparisons easier
    sage["vdW"].switch_width = Quantity(0.0, "nanometer")

    parmed = pytest.importorskip("parmed")

    topology = solvate_topology(
        topology=caffeine.to_topology(),
        nacl_conc=Quantity("0.0 mole/liter"),
        padding=Quantity("1.0 nanometer"),
        box_shape=UNIT_CUBE,
        target_density=Quantity("0.7 gram / centimeter**3"),
    )

    interchange = sage.create_interchange(topology)

    interchange.to_amber("solvated_caffeine")

    structure = parmed.load_file("solvated_caffeine.prmtop")

    # these should all be oxygens
    assert set([a.charge for a in structure.atoms][caffeine.n_atoms :: 3]) == {-0.834}

    # these should all be hydrogens
    assert set([a.charge for a in structure.atoms][caffeine.n_atoms + 1 :: 3]) == {0.417}
    assert set([a.charge for a in structure.atoms][caffeine.n_atoms + 2 :: 3]) == {0.417}

    get_openmm_energies(interchange, combine_nonbonded_forces=False).compare(
        get_amber_energies(interchange),
        tolerances={
            "Angle": Quantity("1e10 kilojoule_per_mole"),  # constraints weirdness
            "vdW": Quantity("0.1 kilojoule_per_mole"),
            "Electrostatics": Quantity("0.5 kilojoule_per_mole"),
        },
    )


@skip_if_missing("openmm")
@pytest.mark.parametrize("valence_handler", ["Bonds", "Angles", "ProperTorsions", "ImproperTorsions"])
def test_issue_1433_openmm(sage, valence_handler):
    """Test that changes in Collection.set_force_field_parameters are reflected in re-computed energies."""
    interchange = sage.create_interchange(
        MoleculeWithConformer.from_smiles(
            "CC(=O)NC",
            allow_undefined_stereo=True,
        ).to_topology(),
    )

    valence_parameters = interchange[valence_handler].get_force_field_parameters()

    original_energy = get_openmm_energies(interchange).total_energy.m

    # multiple force constant (0th column for each handler) by arbtirarily large number
    valence_parameters[:, 0] = valence_parameters[:, 0] * 100

    # impropers do basically nothing energetically, so also shift phase by a few degrees
    # to force an energy change
    if valence_handler == "ImproperTorsions":
        valence_parameters[:, 2] = valence_parameters[:, 2] + 10  # implicit degrees

    interchange[valence_handler].set_force_field_parameters(valence_parameters)

    new_energy = get_openmm_energies(interchange).total_energy.m

    assert abs(original_energy - new_energy) > 0.1


@pytest.mark.parametrize("engine", ["openmm", "amber", "gromacs"])
def test_issue_1450_reassign_vdw(sage, water_dimer, engine):
    """Sanity check that re-assigning vdW parameters actually changes energies."""

    if engine == "amber" and not has_executable("sander"):
        pytest.skip(reason="sander not installed")

    pytest.importorskip("openmm")

    original_interchange = sage.create_interchange(water_dimer)

    match engine:
        case "openmm":
            original_energy = get_openmm_energies(original_interchange, combine_nonbonded_forces=False).energies["vdW"]
        case "amber":
            original_energy = get_amber_energies(original_interchange).energies["vdW"]
        case "gromacs":
            original_energy = get_gromacs_energies(original_interchange).energies["vdW"]

    for potential in original_interchange["vdW"].potentials.values():
        potential.parameters["epsilon"] = potential.parameters["epsilon"] * 100

    match engine:
        case "openmm":
            new_energy = get_openmm_energies(original_interchange, combine_nonbonded_forces=False).energies["vdW"]
        case "amber":
            new_energy = get_amber_energies(original_interchange).energies["vdW"]
        case "gromacs":
            new_energy = get_gromacs_energies(original_interchange).energies["vdW"]

    assert original_energy != new_energy, "Energies should be different after modifying vdW parameters"


@pytest.mark.parametrize("engine", ["openmm", "amber", "gromacs"])
@pytest.mark.parametrize("valence_handler", ["Bonds", "Angles", "ProperTorsions"])
def test_issue_1450_reassign_valence(sage, valence_handler, engine):
    """Sanity check that re-assigning valence parameters actually changes energies."""

    pytest.importorskip("openmm")

    if engine == "amber" and not has_executable("sander"):
        pytest.skip(reason="sander not installed")

    topology = MoleculeWithConformer.from_smiles(
        "CC(=O)NC",
        allow_undefined_stereo=True,
    ).to_topology()

    topology.box_vectors = Quantity(numpy.eye(3) * 10.0, "nanometer")

    original_interchange = sage.create_interchange(topology)

    match valence_handler:
        case "Bonds":
            shorthand = "Bond"
        case "Angles":
            shorthand = "Angle"
        case "ProperTorsions":
            shorthand = "Torsion"

    match engine:
        case "openmm":
            original_energy = get_openmm_energies(original_interchange).energies[shorthand]
        case "amber":
            original_energy = get_amber_energies(original_interchange).energies[shorthand]
        case "gromacs":
            original_energy = get_gromacs_energies(original_interchange).energies[shorthand]

    for potential in original_interchange[valence_handler].potentials.values():
        potential.parameters["k"] = potential.parameters["k"] * 100

        if "phase" in potential.parameters:
            potential.parameters["phase"] = potential.parameters["phase"] + Quantity(20, "degree")

    match engine:
        case "openmm":
            new_energy = get_openmm_energies(original_interchange).energies[shorthand]
        case "amber":
            new_energy = get_amber_energies(original_interchange).energies[shorthand]
        case "gromacs":
            new_energy = get_gromacs_energies(original_interchange).energies[shorthand]

    assert original_energy != new_energy, "Energies should be different after modifying valence parameters"


@skip_if_missing("openmm")
@pytest.mark.parametrize("valence_handler", ["Bonds", "Angles", "ProperTorsions", "ImproperTorsions"])
def test_issue_1234_openmm(sage, valence_handler):
    """Test that modifications to a `ForceField` object are reflected in re-creating new `Interchange`s."""
    topology = MoleculeWithConformer.from_smiles(
        "CC(=O)NC",
        allow_undefined_stereo=True,
    ).to_topology()
    original_interchange = sage.create_interchange(topology)

    original_energy = get_openmm_energies(original_interchange).total_energy.m

    for parameter in sage[valence_handler].parameters:
        if hasattr(parameter, "k"):
            if isinstance(parameter.k, list):
                parameter.k = [k * 100 for k in parameter.k]
            else:
                parameter.k = parameter.k * 100

        # Impropers contribute basically nothing to energies, so modifying the force constant doesn't do much.
        # Instead, shift the phase by 5 degrees, which should have a much more significant effect on energies.
        if hasattr(parameter, "phase"):
            parameter.phase = [phase + Quantity("5 degree") for phase in parameter.phase]
        elif hasattr(parameter, "length"):
            parameter.length = parameter.length * 1.2
        elif hasattr(parameter, "angle"):
            parameter.angle = parameter.angle * 1.1
        else:
            raise ValueError(f"Don't know how to modify parameter with k of type {type(parameter.k)}")

    new_interchange = sage.create_interchange(topology)

    new_energy = get_openmm_energies(new_interchange).total_energy.m

    # energies should be different, since parameters are (wildly!) different
    assert abs(original_energy - new_energy) > 0.1


def test_clear_caches_with_dead_weakref_proxy():
    """Reproduce a bug where _clear_caches raises ReferenceError when dead
    weakref proxies exist in gc.get_objects(). See
    https://github.com/openforcefield/openff-interchange/issues/1453
    """
    import gc
    import weakref

    class _Dummy:
        pass

    # Create a dead weakref proxy that remains tracked by gc
    obj = _Dummy()
    proxy = weakref.proxy(obj)
    container = [proxy]  # noqa: F841
    del obj
    gc.collect()

    molecule = MoleculeWithConformer.from_smiles("C")
    force_field = ForceField("openff-2.2.0.offxml")

    # This raised ReferenceError before the fix
    Interchange.from_smirnoff(force_field=force_field, topology=[molecule])


def test_issue_1461(sage, opc, ethanol):
    with warnings.catch_warnings():
        warnings.simplefilter("error", PresetChargesAndVirtualSitesWarning)

        sage_opc = sage.combine(opc)
        sage_opc.create_interchange(ethanol.to_topology(), charge_from_molecules=[])


@pytest.mark.parametrize("attribute", ["box", "positions"])
def test_issue_1480(sage, ethanol, attribute):
    """Test that positions and box vectors are stored in nm after being set with non-nm length units."""
    out = sage.create_interchange(ethanol.to_topology())

    match attribute:
        case "box":
            out.box = Quantity(numpy.eye(3) * 35, "angstrom")
        case "positions":
            out.positions = Quantity(
                _rng.random(
                    (out.topology.n_atoms, 3),
                ),
                "angstrom",
            )

    assert str(getattr(out, attribute).units) == "nanometer"

    if attribute == "box":
        box = out.box

        assert box[0][0] == box[1][1] == box[2][2]

        assert str(box[0][0].units) == "nanometer"
        assert box[0][0].m == pytest.approx(3.5)
        assert str(box[1][0].units) == "nanometer"
        assert box[1][0].m == pytest.approx(0.0)
