"""Test charge assignment logging."""

import logging
import re
from collections import defaultdict

import pytest
from openff.toolkit import ForceField, Molecule, NAGLToolkitWrapper, RDKitToolkitWrapper, Topology
from openff.toolkit.utils import OPENEYE_AVAILABLE
from openff.toolkit.utils.toolkit_registry import ToolkitRegistry, toolkit_registry_manager

from openff.interchange._tests import get_protein

"""
Hierarchy is
1. Match molecules with preset charges
2. Match chemical environment against library charges
3. Run NAGLCharges (if present in the FF)
3. Run charge method in ChargeIncrementModel section (if present in the FF) and then
   match chemical environment against charge increments
4. Run AM1-BCC (or a variant) on remaining molecules (if present in the FF)

Test cases
----------

Note that each case uses pre-NAGL Sage

0. Sage with a ligand in vacuum
* Ligand gets AM1-BCC

1. Sage with a ligand in water/ions
* Water gets library charges
* Ions get library charges
* Non-water solvent gets AM1-BCC
* Ligands get AM1-BCC

2. Sage with a ligand in non-water solvent
* Non-water solvent gets AM1-BCC
* Ligands get AM1-BCC

2. Sage with a ligand in mixed solvent
* Water gets library charges
* Ions get library charges
* Non-water solvent gets AM1-BCC
* Ligands get AM1-BCC

4. ff14SB with Sage
* Protein gets library charges
* Water gets library charges
* Ions get library charges
* Non-water solvent gets AM1-BCC
* Ligands get AM1-BCC

5. ff14SB with Sage and preset charges on Protein A
* Protein A gets preset charges
* Other proteins get library charges
* Water gets library charges
* Ions get library charges
* Non-water solvent gets AM1-BCC
* Ligands get AM1-BCC

6. Sage with ligand and OPC water
* Water gets library charges
* Water virtual sites get ??
* Ions get library charges
* Non-water solvent gets AM1-BCC
* Ligands get AM1-BCC

7. Sage with preset charges on ligand A
* Water gets library charges
* Ions get library charges
* Non-water solvent gets AM1-BCC
* Ligand A gets preset charges
* Other ligands get AM1-BCC

8. Sage with preset charges on water
* Water gets preset charges
* Ions get library charges
* Non-water solvent gets AM1-BCC
* Ligands get AM1-BCC

9. Sage with (ligand) virtual site parameters
* Water gets preset charges
* Ions get library charges
* Non-water solvent gets AM1-BCC
* Ligand heavy atoms get AM1-BCC and charge increments
* Virtual sites get charge increments

10. AM1-with-custom-BCCs Sage with ligand and ions water
* Water gets library charges
* Ions get library charges
* Ligand gets charge increments

11. Sage with NAGL and ligand in vacuum
* Ligand gets NAGL charges

12. Sage with NAGL and water molecules (library precedence)
* Water gets library charges (precedence over NAGL)

13. Sage with NAGL mixed topology (ligand and water)
* Ligand gets NAGL charges
* Water gets library charges

14. Sage with NAGL and preset charges on ligand
* Ligand gets preset charges (precedence over NAGL)

Other details
* Specifics of charge method (AM1-BCC, AM1-BCC-ELF10, AM1-BCC via NAGL)
* Molecules with preset charges can be similar but not exact enough
* Preset charges and virtual sites is underspecified and cannot be tested
"""

# TODO: Revisit test cases given that NAGL charges are default (2.3.0+, 3.0.0, etc.)
AM1BCC_KEY = "am1bccelf10" if OPENEYE_AVAILABLE else "am1bcc"
NAGL_KEY = "openff-gnn-am1bcc-1.0.0.pt"

INCHI_KEYS = {
    "ethanol": "LFQSCWFLJHTTHZ-UHFFFAOYSA-N",
    "water": "XLYOFNOQVPJJNP-UHFFFAOYSA-N",
    "sodium": "FKNQFGJONOIPTF-UHFFFAOYSA-N",
    "chloride": "VEXZGXHMUGYJMC-UHFFFAOYSA-M",
    "methylamine": "BAVYZALUXZFZLV-UHFFFAOYSA-N",
    "protein": "QIVKPXNXCAPCCO-WDSKDSINSA-N",
}


def map_methods_to_inchikey(caplog: pytest.LogCaptureFixture) -> dict[str, list[int]]:
    """
    Map partial charge assignment methods to (sorted) molecule InChI.
    """
    debug = defaultdict(list)

    for record in caplog.records:
        if record.name.startswith("openff.interchange"):
            # This is a specific one-off logging event not relevant to this function
            if "Could not generate InChIKey for molecule with Hill formula" in record.msg:
                continue
            assert record.levelname == "DEBUG", "Only DEBUG logs are expected."
        else:
            # skip logged warnings from upstreams/other packages
            continue

        message = record.msg

        if message.startswith("Charge section LibraryCharges"):
            _, molecule = message.split(", ")

            debug["library"].append(molecule.split("InChIKey ")[-1])

        elif message.startswith("Charge section ToolkitAM1BCC"):
            _, method, molecule = message.split(", ")

            assert method.split()[-1] == AM1BCC_KEY, f"Expected method {AM1BCC_KEY} but got {method.split()[-1]}"

            debug[method.split()[-1]].append(molecule.split("InChIKey ")[-1])

        elif message.startswith("Charge section NAGLCharges"):
            _, _, molecule = message.split(", ")

            debug["NAGL"].append(molecule.split("InChIKey ")[-1])

        # without also pulling the virtual site - particle mapping (which is different for each engine)
        # it's hard to store more debugrmation than the orientation atoms that are affected by each
        # virtual site's charges
        elif message.startswith("Charge section VirtualSites"):
            orientation_atoms: list[int] = [
                int(value.strip()) for value in re.findall(r"\((.*?)\)", message)[0].split(",")
            ]

            for atom in orientation_atoms:
                debug["orientation"].append(atom)

        elif message.startswith("Preset charges"):
            debug["preset"].append(message.split("InChIKey ")[-1])

        elif message.startswith("Charge section ChargeIncrementModel"):
            if "using charge method" in message:
                _, method, molecule = message.split(", ")
                debug[f"chargeincrements_{method.split()[-1]}"].append(
                    molecule.split("InChIKey ")[-1],
                )

            elif "applying charge increment" in message:
                # TODO: Store the "other" atoms?
                debug["chargeincrements"].append(int(message.split("atom ")[1].split(" ")[0]))

        else:
            raise ValueError(f"Unexpected log message {message}")

    return {key: sorted(val) for key, val in debug.items()}


def _ensure_pre_nagl_sage(sage: ForceField) -> ForceField:
    """Revert a verison of Sage to pre-NAGL by swapping out NAGLChharges handler for ToolkitAM1BCC."""
    try:
        sage.deregister_parameter_handler("NAGLCharges")
        sage.get_parameter_handler("ToolkitAM1BCC")

        return sage

    except KeyError:
        return sage


@pytest.fixture
def sage_with_nagl_chargeincrements(sage_no_nagl):
    from openff.toolkit.typing.engines.smirnoff.parameters import ChargeIncrementModelHandler, ChargeIncrementType

    sage_no_nagl.register_parameter_handler(
        parameter_handler=ChargeIncrementModelHandler(
            version=0.3,
            partial_charge_method=NAGL_KEY,
        ),
    )

    # Add dummy "BCCs" for testing, even though this model has BCCs built into
    # the partial charge assignment method
    sage_no_nagl["ChargeIncrementModel"].add_parameter(
        parameter=ChargeIncrementType(
            smirks="[#6:1]-[#1:2]",
            charge_increment=[
                "0.1 elementary_charge",
                "-0.1 elementary_charge",
            ],
        ),
    )

    return sage_no_nagl


@pytest.fixture
def ions() -> Topology:
    return Topology.from_molecules(
        [
            Molecule.from_smiles(smiles)
            for smiles in [
                "[Na+]",
                "[Cl-]",
                "[Na+]",
                "[Cl-]",
            ]
        ],
    )


@pytest.fixture
def ligand() -> Molecule:
    return Molecule.from_mapped_smiles("[C:1]([C:2]([O:3][H:9])([H:7])[H:8])([H:4])([H:5])[H:6]")


@pytest.fixture
def solvent() -> Molecule:
    return Molecule.from_mapped_smiles("[H:3][C:1]([H:4])([H:5])[N:2]([H:6])[H:7]")


def ligand_in_solvent(ligand, solvent) -> Topology:
    return Topology.from_molecules(
        [
            ligand,
            solvent,
            solvent,
            solvent,
        ],
    )


@pytest.fixture
def water_and_ions(water, ions) -> Topology:
    topology = Topology.from_molecules(3 * [water])
    for molecule in ions.molecules:
        topology.add_molecule(molecule)

    return topology


@pytest.fixture
def ligand_and_water_and_ions(ligand, water_and_ions) -> Topology:
    topology = ligand.to_topology()

    for molecule in water_and_ions.molecules:
        topology.add_molecule(molecule)

    return topology


"""
0.xSage with just a ligand
1.xSage with a ligand in water/ions
2.xSage with a ligand in non-water solvent
3.xSage with a ligand in mixed solvent
4.xff14SB with Sage
5.xff14SB with Sage and preset charges on Protein A
6.xSage with ligand and OPC water
7.xSage with preset charges on ligand A
8.xSage with preset charges on water
9.xSage with (ligand) virtual site parameters
10.xAM1-with-custom-BCCs Sage with ligand and ions water
11.xSage with NAGL and ligand in vacuum
12.xSage with NAGL and water molecules (library precedence)
13.xSage with NAGL mixed topology (ligand and water)
14.xSage with NAGL and preset charges on ligand
"""


def test_case0(caplog, sage_no_nagl, ligand):
    with caplog.at_level(logging.DEBUG, logger="openff.interchange"):
        sage_no_nagl.create_interchange(ligand.to_topology())

        debug = map_methods_to_inchikey(caplog)

        # atoms 0 through 8 are ethanol, getting AM1-BCC
        assert debug[AM1BCC_KEY] == [INCHI_KEYS["ethanol"]]


def test_case1(caplog, sage_no_nagl, ligand_and_water_and_ions):
    with caplog.at_level(logging.DEBUG, logger="openff.interchange"):
        sage_no_nagl.create_interchange(ligand_and_water_and_ions)

        debug = map_methods_to_inchikey(caplog)

        # atoms 0 through 8 are ethanol, getting AM1-BCC
        assert debug[AM1BCC_KEY] == [INCHI_KEYS["ethanol"]]

        # atoms 9 through 21 are water + ions, getting library charges
        assert debug["library"] == [
            INCHI_KEYS["sodium"],
            INCHI_KEYS["chloride"],
            INCHI_KEYS["water"],
        ]


def test_case2(caplog, sage_no_nagl, ligand, solvent):
    topology = Topology.from_molecules([ligand, solvent, solvent, solvent])

    with caplog.at_level(logging.DEBUG, logger="openff.interchange"):
        sage_no_nagl.create_interchange(topology)

        debug = map_methods_to_inchikey(caplog)

        # everything should get AM1-BCC charges
        assert debug[AM1BCC_KEY] == [INCHI_KEYS["methylamine"], INCHI_KEYS["ethanol"]]


def test_case3(caplog, sage_no_nagl, ligand_and_water_and_ions, solvent):
    for index in range(3):
        ligand_and_water_and_ions.add_molecule(solvent)

    ligand_and_water_and_ions.molecule(0).assign_partial_charges(partial_charge_method="gasteiger")

    with caplog.at_level(logging.DEBUG, logger="openff.interchange"):
        sage_no_nagl.create_interchange(
            ligand_and_water_and_ions,
        )

        debug = map_methods_to_inchikey(caplog)

        # atoms 0 through 8 are ethanol, getting AM1-BCC,
        # and also solvent molecules (starting at index 22)
        assert debug[AM1BCC_KEY] == [INCHI_KEYS["methylamine"], INCHI_KEYS["ethanol"]]

        # atoms 9 through 21 are water + ions, getting library charges
        assert debug["library"] == [
            INCHI_KEYS["sodium"],
            INCHI_KEYS["chloride"],
            INCHI_KEYS["water"],
        ]


@pytest.mark.slow
@pytest.mark.parametrize("preset_on_protein", [True, False])
def test_cases4_5(caplog, ligand_and_water_and_ions, preset_on_protein):
    ff = ForceField("ff14sb_off_impropers_0.0.3.offxml", "openff-2.0.0.offxml")

    complex = get_protein("MainChain_ALA_ALA").to_topology()

    for molecule in ligand_and_water_and_ions.molecules:
        complex.add_molecule(molecule)

    if preset_on_protein:
        complex.molecule(0).assign_partial_charges(partial_charge_method="zeros")

    with caplog.at_level(logging.DEBUG, logger="openff.interchange"):
        if preset_on_protein:
            ff.create_interchange(complex, charge_from_molecules=[complex.molecule(0)])
        else:
            ff.create_interchange(complex)

        debug = map_methods_to_inchikey(caplog)

        assert debug[AM1BCC_KEY] == [INCHI_KEYS["ethanol"]]

        if preset_on_protein:
            # protein gets preset charges
            assert debug["preset"] == [
                INCHI_KEYS["protein"],
            ]

            # everything after the protein and ligand should get library charges
            assert debug["library"] == [
                INCHI_KEYS["sodium"],
                INCHI_KEYS["chloride"],
                INCHI_KEYS["water"],
            ]
        else:
            # the protein and everything after the ligand should get library charges
            assert debug["library"] == [
                INCHI_KEYS["sodium"],
                INCHI_KEYS["protein"],
                INCHI_KEYS["chloride"],
                INCHI_KEYS["water"],
            ]


def test_case6(caplog, ligand, water):
    force_field = ForceField("openff-2.0.0.offxml", "opc.offxml")

    topology = Topology.from_molecules([ligand, water, water, water])

    with caplog.at_level(logging.DEBUG, logger="openff.interchange"):
        force_field.create_interchange(topology)

        debug = map_methods_to_inchikey(caplog)

        # atoms 0 through 8 are ethanol, getting AM1-BCC charges
        assert debug[AM1BCC_KEY] == [INCHI_KEYS["ethanol"]]

        # atoms 9 through 17 are water atoms, getting library charges
        assert debug["library"] == [INCHI_KEYS["water"]]

        # particles 18 through 20 are water virtual sites, but the current logging strategy
        # makes it hard to match these up (and the particle indices are different OpenMM/GROMACS/etc)

        # can still check that orientation atoms are subject to virtual site
        # charge increments (even if the increment is +0.0 e)
        assert debug["orientation"] == [*range(9, 18)]


def test_case7(caplog, sage_no_nagl, ligand_and_water_and_ions):
    ligand_and_water_and_ions.molecule(0).assign_partial_charges(partial_charge_method="gasteiger")

    with caplog.at_level(logging.DEBUG, logger="openff.interchange"):
        sage_no_nagl.create_interchange(
            ligand_and_water_and_ions,
            charge_from_molecules=[ligand_and_water_and_ions.molecule(0)],
        )

        debug = map_methods_to_inchikey(caplog)

        # atoms 0 through 8 are ethanol, getting preset charges
        assert debug["preset"] == [INCHI_KEYS["ethanol"]]

        # atoms 9 through 21 are water + ions, getting library charges
        assert debug["library"] == [
            INCHI_KEYS["sodium"],
            INCHI_KEYS["chloride"],
            INCHI_KEYS["water"],
        ]


def test_case8(caplog, sage_no_nagl, water_and_ions):
    water_and_ions.molecule(0).assign_partial_charges(partial_charge_method="gasteiger")

    with caplog.at_level(logging.DEBUG, logger="openff.interchange"):
        sage_no_nagl.create_interchange(
            water_and_ions,
            charge_from_molecules=[water_and_ions.molecule(0)],
        )

        debug = map_methods_to_inchikey(caplog)

        # atoms 0 through 8 are water, getting preset charges
        assert debug["preset"] == [INCHI_KEYS["water"]]

        # atoms 9 through 12 are ions, getting library charges
        assert debug["library"] == [INCHI_KEYS["sodium"], INCHI_KEYS["chloride"]]


def test_case9(caplog, sage_with_bond_charge):
    with caplog.at_level(logging.DEBUG, logger="openff.interchange"):
        _ensure_pre_nagl_sage(sage_with_bond_charge).create_interchange(
            Molecule.from_mapped_smiles(
                "[H:3][C:1]([H:4])([H:5])[Cl:2]",
            ).to_topology(),
        )

        debug = map_methods_to_inchikey(caplog)

        # atoms 0 through 5 are ligand, getting AM1-BCC charges
        assert debug[AM1BCC_KEY] == ["NEHMKBQYUWJMIP-UHFFFAOYSA-N"]

        # atoms 0 and 1 are the orientation atoms of the sigma hole virtual site
        assert debug["orientation"] == [0, 1]


def test_case10(caplog, sage_with_nagl_chargeincrements, ligand):
    pytest.importorskip("openff.nagl")
    pytest.importorskip("rdkit")

    with caplog.at_level(logging.DEBUG, logger="openff.interchange"):
        with toolkit_registry_manager(
            toolkit_registry=ToolkitRegistry(
                toolkit_precedence=[NAGLToolkitWrapper, RDKitToolkitWrapper],
            ),
        ):
            _ensure_pre_nagl_sage(sage_with_nagl_chargeincrements).create_interchange(
                ligand.to_topology(),
            )

        debug = map_methods_to_inchikey(caplog)

        # atoms 0 through 8 are ligand, getting NAGL charges
        assert debug[f"chargeincrements_{NAGL_KEY}"] == ["LFQSCWFLJHTTHZ-UHFFFAOYSA-N"]

        # TODO: These are logged symmetrically (i.e. hydrogens are listed)
        # even though the charges appear to be correct, assert should
        # simply by == [0, 1] since the hydrogens shouldn't be caught
        assert 0 in debug["chargeincrements"]
        assert 1 in debug["chargeincrements"]

        # the standard AM1-BCC should not have ran
        assert AM1BCC_KEY not in debug


def test_case11(caplog, sage, ligand):
    """Test that NAGL charge assignment is properly logged."""
    pytest.importorskip("openff.nagl")

    with caplog.at_level(logging.DEBUG, logger="openff.interchange"):
        sage.create_interchange(ligand.to_topology())

        debug = map_methods_to_inchikey(caplog)

        # Should log NAGL charges for all atoms
        assert "NAGL" in debug
        assert debug["NAGL"] == [INCHI_KEYS["ethanol"]]


def test_case12(caplog, sage, water):
    """Test logging when LibraryCharges takes precedence over NAGL."""
    pytest.importorskip("openff.nagl")

    topology = Topology.from_molecules([water, water])

    with caplog.at_level(logging.DEBUG, logger="openff.interchange"):
        sage.create_interchange(topology)

        debug = map_methods_to_inchikey(caplog)

        # Water should get library charges, not NAGL
        assert debug["library"] == [INCHI_KEYS["water"]]  # 2 water molecules
        assert "NAGL" not in debug


def test_case13(caplog, sage, ligand, water):
    """Test logging with mixed molecule types (some library, some NAGL)."""
    pytest.importorskip("openff.nagl")

    topology = Topology.from_molecules([ligand, water])

    with caplog.at_level(logging.DEBUG, logger="openff.interchange"):
        sage.create_interchange(topology)

        debug = map_methods_to_inchikey(caplog)

        # Ligand should get NAGL charges
        assert debug["NAGL"] == [INCHI_KEYS["ethanol"]]

        # Water should get library charges
        assert debug["library"] == [INCHI_KEYS["water"]]


def test_case14(caplog, sage, ligand):
    """Test logging when preset charges are used instead of NAGL."""
    pytest.importorskip("openff.nagl")

    ligand.assign_partial_charges("gasteiger")

    with caplog.at_level(logging.DEBUG, logger="openff.interchange"):
        sage.create_interchange(
            ligand.to_topology(),
            charge_from_molecules=[ligand],
        )

        debug = map_methods_to_inchikey(caplog)

        # Should use preset charges, not NAGL
        assert debug["preset"] == [INCHI_KEYS["ethanol"]]
        assert "NAGL" not in debug


def test_inchi_fallback(caplog, sage):
    """Test that molecules that fail InChIKey generation are still logged in some way."""
    from openff.toolkit.utils.toolkits import OpenEyeToolkitWrapper

    # TODO: Might be a toolkit bug that needs to be worked around
    with caplog.at_level(logging.DEBUG, logger="openff.interchange"):
        sage.create_interchange(
            Molecule.from_smiles(342 * "C").to_topology(),
        )

        debug = map_methods_to_inchikey(caplog)

        if OpenEyeToolkitWrapper.is_available():
            assert debug["NAGL"] == ["UNKNOWN_INCHIKEY"]
        else:
            # but RDKit can generate an InChIKey
            assert debug["NAGL"] == ["FGIJOQKVOYCTPJ-UHFFFAOYSA-N"]
