# Copyright 2026 AlQuraishi Laboratory
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import networkx as nx
import numpy as np
import pytest
from biotite.structure import Atom, AtomArray, BondList, BondType

from openfold3.core.data.primitives.structure.cleanup import (
    filter_fully_atomized_bonds,
    prefilter_bonds,
    remove_covalent_nonprotein_chains,
)
from openfold3.core.data.primitives.structure.component import find_cross_chain_bonds
from openfold3.core.data.resources.residues import MoleculeType
from openfold3.tests.custom_assert_utils import assert_atomarray_equal
from openfold3.tests.data_utils import create_atomarray_with_bondlist

# This helps conciseness to avoid Ruff line-break
PROTEIN = MoleculeType.PROTEIN
LIGAND = MoleculeType.LIGAND
DNA = MoleculeType.DNA


# -- Test Case Data ---

# Atom array that covers all the following properties:
# - intra-chain dative bond
# - inter-chain dative bond
# - intra-chain consecutive polymer bonds
# - intra-chain nonconsecutive polymer bonds
# - inter-chain polymer bonds
# - 2 long bonds (longer than 2.4 Å)

# atom_name is only set for user interpretability and irrelevant for the test

# fmt: off
all_atoms=[
    # Protein chain A
    Atom([0, 0, 0], chain_id="A", res_id=1, atom_name="backbone", molecule_type_id=PROTEIN), # noqa: E501
    Atom([0, 0, 0], chain_id="A", res_id=1, atom_name="sidechain", molecule_type_id=PROTEIN), # noqa: E501
    Atom([0, 0, 0], chain_id="A", res_id=2, atom_name="backbone", molecule_type_id=PROTEIN), # noqa: E501
    Atom([0, 0, 0], chain_id="A", res_id=2, atom_name="sidechain", molecule_type_id=PROTEIN), # noqa: E501
    Atom([0, 0, 0], chain_id="A", res_id=3, atom_name="backbone", molecule_type_id=PROTEIN), # noqa: E501
    Atom([0, 0, 0], chain_id="A", res_id=3, atom_name="sidechain", molecule_type_id=PROTEIN), # noqa: E501

    # Ligand chain B
    Atom([0, 0, 0], chain_id="B", res_id=1, atom_name="ligand", molecule_type_id=LIGAND), # noqa: E501
    Atom([0, 0, 0], chain_id="B", res_id=1, atom_name="ligand", molecule_type_id=LIGAND), # noqa: E501
    Atom([0, 0, 2.41], chain_id="B", res_id=1, atom_name="ligand", molecule_type_id=LIGAND), # noqa: E501

    # DNA chain C
    Atom([2.41, 0, 0], chain_id="C", res_id=1, atom_name="backbone", molecule_type_id=DNA), # noqa: E501
    Atom([2.41, 0, 0], chain_id="C", res_id=1, atom_name="sidechain", molecule_type_id=DNA), # noqa: E501
    Atom([0, 0, 0], chain_id="C", res_id=2, atom_name="backbone", molecule_type_id=DNA), # noqa: E501
    Atom([0, 0, 0], chain_id="C", res_id=2, atom_name="sidechain", molecule_type_id=DNA), # noqa: E501
    Atom([0, 0, 0], chain_id="C", res_id=3, atom_name="backbone", molecule_type_id=DNA), # noqa: E501
    Atom([0, 0, 0], chain_id="C", res_id=3, atom_name="sidechain", molecule_type_id=DNA), # noqa: E501
    Atom([0, 0, 0], chain_id="C", res_id=4, atom_name="backbone", molecule_type_id=DNA), # noqa: E501
    Atom([0, 0, 0], chain_id="C", res_id=4, atom_name="sidechain", molecule_type_id=DNA), # noqa: E501
    Atom([0, 0, 0], chain_id="C", res_id=5, atom_name="backbone", molecule_type_id=DNA), # noqa: E501
    Atom([0, 0, 0], chain_id="C", res_id=5, atom_name="sidechain", molecule_type_id=DNA), # noqa: E501
]
# fmt: on

# - References for special bonds -

# Chain B
intra_chain_long_dative = (7, 8, BondType.COORDINATION)

# Chain A to Chain B (dative)
inter_chain_dative = (3, 6, BondType.COORDINATION)

# Chain A to Chain B (standard) <- should not be removed
inter_chain_standard = (2, 7, BondType.SINGLE)

# Chain C
intra_chain_poly_link = (14, 18, BondType.SINGLE)

# Chain A to Chain C
inter_chain_poly_link = (5, 12, BondType.SINGLE)

# Chain C
long_bond_2 = (9, 11, BondType.SINGLE)

long_bond_set = {intra_chain_long_dative, long_bond_2}

# - Bond set -
bond_set = set(
    (
        # -- Intra-chain standard bonds --
        # Chain A
        (0, 1, BondType.SINGLE),
        (0, 2, BondType.SINGLE),
        (2, 3, BondType.SINGLE),
        (2, 4, BondType.SINGLE),
        (4, 5, BondType.SINGLE),
        # Chain B
        (6, 7, BondType.SINGLE),
        # Chain C
        (9, 10, BondType.SINGLE),
        long_bond_2,
        (11, 12, BondType.SINGLE),
        (11, 13, BondType.SINGLE),
        (13, 14, BondType.SINGLE),
        (13, 15, BondType.SINGLE),
        (15, 16, BondType.SINGLE),
        (15, 17, BondType.SINGLE),
        (17, 18, BondType.SINGLE),
        # -- Intra-chain dative bond --
        intra_chain_long_dative,
        # -- Inter-chain dative bond --
        inter_chain_dative,
        # -- Inter-chain standard bond --
        inter_chain_standard,
        # -- Intra-chain nonconsecutive polymer crosslink --
        intra_chain_poly_link,
        # -- Inter-chain polymer crosslink --
        inter_chain_poly_link,
    )
)

# Sorting the bonds here is important for unit-test equivalence
atom_array_filter_bonds = create_atomarray_with_bondlist(
    all_atoms, np.array(sorted(bond_set))
)

# -- Test Cases --

# - Case 1: all active -
# remove_inter_chain_dative: True
# remove_inter_chain_poly_links: True
# remove_intra_chain_poly_links: True
# remove_longer_than: 2.4
case_all_true_bond_set = (
    bond_set
    - {inter_chain_dative, inter_chain_poly_link, intra_chain_poly_link}
    - long_bond_set
)
case_all_true_bond_array = np.array(sorted(case_all_true_bond_set))

# - Case 2: all inactive -
# remove_inter_chain_dative: False
# remove_inter_chain_poly_links: False
# remove_intra_chain_poly_links: False
# remove_longer_than: None
case_all_false_bond_set = bond_set
case_all_false_bond_array = np.array(sorted(case_all_false_bond_set))

# - Case 3: only inter-chain dative active -
# remove_inter_chain_dative: True
# remove_inter_chain_poly_links: False
# remove_intra_chain_poly_links: False
# remove_longer_than: None
case_only_inter_chain_dative = bond_set - {inter_chain_dative}
case_only_inter_chain_dative_bond_array = np.array(sorted(case_only_inter_chain_dative))

# - Case 4: only inter-chain poly links active -
# remove_inter_chain_dative: False
# remove_inter_chain_poly_links: True
# remove_intra_chain_poly_links: False
# remove_longer_than: None
case_only_inter_chain_poly_links = bond_set - {inter_chain_poly_link}
case_only_inter_chain_poly_links_bond_array = np.array(
    sorted(case_only_inter_chain_poly_links)
)

# - Case 5: only intra-chain poly links active -
# remove_inter_chain_dative: False
# remove_inter_chain_poly_links: False
# remove_intra_chain_poly_links: True
# remove_longer_than: None
case_only_intra_chain_poly_links = bond_set - {intra_chain_poly_link}
case_only_intra_chain_poly_links_bond_array = np.array(
    sorted(case_only_intra_chain_poly_links)
)

# - Case 6: only longer_than active -
# remove_inter_chain_dative: False
# remove_inter_chain_poly_links: False
# remove_intra_chain_poly_links: False
# remove_longer_than: 2.4
case_only_longer_than = bond_set - long_bond_set
case_only_longer_than_bond_array = np.array(sorted(case_only_longer_than))


@pytest.mark.parametrize(
    "params",
    [
        # All True
        {
            "atom_array": atom_array_filter_bonds,
            "remove_inter_chain_dative": True,
            "remove_inter_chain_poly_links": True,
            "remove_intra_chain_poly_links": True,
            "remove_longer_than": 2.4,
            "expected_bondlist": case_all_true_bond_array,
        },
        # All False
        {
            "atom_array": atom_array_filter_bonds,
            "remove_inter_chain_dative": False,
            "remove_inter_chain_poly_links": False,
            "remove_intra_chain_poly_links": False,
            "remove_longer_than": None,
            "expected_bondlist": case_all_false_bond_array,
        },
        # Only remove inter-chain dative
        {
            "atom_array": atom_array_filter_bonds,
            "remove_inter_chain_dative": True,
            "remove_inter_chain_poly_links": False,
            "remove_intra_chain_poly_links": False,
            "remove_longer_than": None,
            "expected_bondlist": case_only_inter_chain_dative_bond_array,
        },
        # Only remove inter-chain polymer links
        {
            "atom_array": atom_array_filter_bonds,
            "remove_inter_chain_dative": False,
            "remove_inter_chain_poly_links": True,
            "remove_intra_chain_poly_links": False,
            "remove_longer_than": None,
            "expected_bondlist": case_only_inter_chain_poly_links_bond_array,
        },
        # Only remove intra-chain polymer links
        {
            "atom_array": atom_array_filter_bonds,
            "remove_inter_chain_dative": False,
            "remove_inter_chain_poly_links": False,
            "remove_intra_chain_poly_links": True,
            "remove_longer_than": None,
            "expected_bondlist": case_only_intra_chain_poly_links_bond_array,
        },
        # Only remove longer than 2.4
        {
            "atom_array": atom_array_filter_bonds,
            "remove_inter_chain_dative": False,
            "remove_inter_chain_poly_links": False,
            "remove_intra_chain_poly_links": False,
            "remove_longer_than": 2.4,
            "expected_bondlist": case_only_longer_than_bond_array,
        },
    ],
    ids=[
        "all_true",
        "all_false",
        "only_rm_inter_chain_dative",
        "only_rm_inter_chain_poly_links",
        "only_rm_intra_chain_poly_links",
        "only_rm_longer_than_2.4",
    ],
)
def test_prefilter_bonds(params):
    """Tests whether the bond prefiltering works as expected."""

    # Create a copy of the atom array to avoid modifying the original
    atom_array = params["atom_array"]
    expected_atom_array = atom_array.copy()
    expected_bondlist = params["expected_bondlist"]

    # Set the expected bond list
    expected_atom_array.bonds = BondList(len(expected_atom_array), expected_bondlist)

    atom_array_filtered = prefilter_bonds(
        atom_array=atom_array,
        remove_inter_chain_dative=params["remove_inter_chain_dative"],
        remove_inter_chain_poly_links=params["remove_inter_chain_poly_links"],
        remove_intra_chain_poly_links=params["remove_intra_chain_poly_links"],
        remove_longer_than=params["remove_longer_than"],
    )

    assert_atomarray_equal(
        atom_array_filtered,
        expected_atom_array,
    )


# -- Test Case Data ---

# - Case 1: no atomized atoms at all -
atom_array_no_atomized = create_atomarray_with_bondlist(
    atoms=[
        Atom([0, 0, 0], chain_id="A", res_id=1, is_atomized=False),
        Atom([0, 0, 0], chain_id="A", res_id=1, is_atomized=False),
        Atom([0, 0, 0], chain_id="A", res_id=1, is_atomized=False),
        Atom([0, 0, 0], chain_id="A", res_id=2, is_atomized=False),
        Atom([0, 0, 0], chain_id="A", res_id=2, is_atomized=False),
        Atom([0, 0, 0], chain_id="A", res_id=3, is_atomized=False),
        Atom([0, 0, 0], chain_id="A", res_id=3, is_atomized=False),
    ],
    bondlist=np.array(
        sorted(
            [
                (0, 1, BondType.SINGLE),
                (1, 2, BondType.SINGLE),
                (2, 3, BondType.SINGLE),
                (3, 4, BondType.SINGLE),
                (4, 5, BondType.SINGLE),
                (5, 6, BondType.SINGLE),
            ]
        )
    ),
)
expected_bondlist_no_atomized = np.array([], dtype=int).reshape(0, 3)

# - Case 2: non-atomized residue binding to atomized ligand -
atom_array_some_atomized = create_atomarray_with_bondlist(
    atoms=[
        Atom([0, 0, 0], chain_id="A", res_id=1, is_atomized=False),
        Atom([0, 0, 0], chain_id="A", res_id=1, is_atomized=False),
        Atom([0, 0, 0], chain_id="B", res_id=1, is_atomized=True),
        Atom([0, 0, 0], chain_id="B", res_id=1, is_atomized=True),
        Atom([0, 0, 0], chain_id="B", res_id=1, is_atomized=True),
    ],
    bondlist=np.array(
        sorted(
            [
                (0, 1, BondType.SINGLE),
                (1, 2, BondType.SINGLE),
                (2, 3, BondType.SINGLE),
                (3, 4, BondType.SINGLE),
            ]
        )
    ),
)
expected_bondlist_some_atomized = np.array(
    [
        (2, 3, BondType.SINGLE),
        (3, 4, BondType.SINGLE),
    ],
    dtype=int,
)

# - Case 3: all atoms are atomized -

# make bondlist fully connected
G = nx.complete_graph(8)
edges = np.array(G.edges())
bondlist_all_atomized = np.hstack(
    (edges, np.full((edges.shape[0], 1), BondType.SINGLE, dtype=int))
)

atom_array_all_atomized = create_atomarray_with_bondlist(
    atoms=[
        Atom([0, 0, 0], chain_id="A", res_id=1, is_atomized=True),
        Atom([0, 0, 0], chain_id="A", res_id=1, is_atomized=True),
        Atom([0, 0, 0], chain_id="A", res_id=1, is_atomized=True),
        Atom([0, 0, 0], chain_id="A", res_id=2, is_atomized=True),
        Atom([0, 0, 0], chain_id="A", res_id=2, is_atomized=True),
        Atom([0, 0, 0], chain_id="B", res_id=1, is_atomized=True),
        Atom([0, 0, 0], chain_id="B", res_id=1, is_atomized=True),
        Atom([0, 0, 0], chain_id="B", res_id=2, is_atomized=True),
    ],
    bondlist=bondlist_all_atomized,
)


@pytest.mark.parametrize(
    ["atom_array", "expected_bondlist"],
    [
        (atom_array_no_atomized, expected_bondlist_no_atomized),
        (atom_array_some_atomized, expected_bondlist_some_atomized),
        (atom_array_all_atomized, bondlist_all_atomized),
    ],
    ids=["no_atomized_bonds", "some_atomized_bonds", "all_atomized_bonds"],
)
def test_filter_fully_atomized_bonds(
    atom_array: AtomArray, expected_bondlist: BondList | np.ndarray
):
    """Tests whether filtering of fully atomized bonds works as expected."""

    # Create a copy of the atom array to avoid modifying the original
    atom_array_expected = atom_array.copy()

    # Set the expected bond list
    atom_array_expected.bonds = BondList(len(atom_array_expected), expected_bondlist)

    # Apply the filter
    filtered_bondlist = filter_fully_atomized_bonds(
        atom_array=atom_array,
    )

    # Assert that the filtered bond list matches the expected bond list
    assert_atomarray_equal(
        filtered_bondlist,
        atom_array_expected,
    )


# -- Test find_cross_chain_bonds --


def test_find_cross_chain_bonds_single_cross_chain():
    """Tests that find_cross_chain_bonds correctly identifies cross-chain bonds."""
    # Minimal protein-ligand complex: 2 atoms in chain A, 2 atoms in chain B
    atoms = [
        Atom([0, 0, 0], chain_id="A", res_id=1),  # idx 0
        Atom([1, 0, 0], chain_id="A", res_id=1),  # idx 1
        Atom([2, 0, 0], chain_id="B", res_id=1),  # idx 2
        Atom([3, 0, 0], chain_id="B", res_id=1),  # idx 3
    ]
    bonds = np.array(
        [
            (0, 1, BondType.SINGLE),  # intra-chain A
            (2, 3, BondType.SINGLE),  # intra-chain B
            (1, 2, BondType.SINGLE),  # cross-chain A->B
        ]
    )
    atom_array = create_atomarray_with_bondlist(atoms, bonds)

    cross_chain_bonds = find_cross_chain_bonds(atom_array)

    assert cross_chain_bonds.shape == (1, 3)
    assert tuple(cross_chain_bonds[0]) == (1, 2, BondType.SINGLE)


def test_find_cross_chain_bonds_no_cross_chain():
    """Tests find_cross_chain_bonds with no cross-chain bonds."""
    # Two separate chains with no bonds between them
    atoms = [
        Atom([0, 0, 0], chain_id="A", res_id=1),
        Atom([1, 0, 0], chain_id="A", res_id=1),
        Atom([10, 0, 0], chain_id="B", res_id=1),
        Atom([11, 0, 0], chain_id="B", res_id=1),
    ]
    bonds = np.array(
        [
            (0, 1, BondType.SINGLE),  # intra-chain A only
            (2, 3, BondType.SINGLE),  # intra-chain B only
        ]
    )
    atom_array = create_atomarray_with_bondlist(atoms, bonds)

    cross_chain_bonds = find_cross_chain_bonds(atom_array)

    assert cross_chain_bonds.shape[0] == 0


def test_find_cross_chain_bonds_includes_coordination():
    """Tests that find_cross_chain_bonds includes coordination (dative) bonds."""
    # Metal ion coordinated to protein residue across chains
    atoms = [
        Atom([0, 0, 0], chain_id="A", res_id=1),  # idx 0: protein atom
        Atom([2.3, 0, 0], chain_id="B", res_id=1),  # idx 1: metal ion
    ]
    bonds = np.array(
        [
            (0, 1, BondType.COORDINATION),  # cross-chain coordination bond
        ]
    )
    atom_array = create_atomarray_with_bondlist(atoms, bonds)

    cross_chain_bonds = find_cross_chain_bonds(atom_array)

    assert cross_chain_bonds.shape == (1, 3)
    assert tuple(cross_chain_bonds[0]) == (0, 1, BondType.COORDINATION)


# -- Test remove_covalent_nonprotein_chains --


def test_remove_covalent_nonprotein_chains_removes_ligand():
    """Tests that ligand chains covalently bonded to protein are removed."""
    # Protein chain A covalently bonded to ligand chain B
    atoms = [
        Atom([0, 0, 0], chain_id="A", res_id=1, molecule_type_id=PROTEIN),
        Atom([1, 0, 0], chain_id="A", res_id=1, molecule_type_id=PROTEIN),
        Atom([2, 0, 0], chain_id="B", res_id=1, molecule_type_id=LIGAND),
        Atom([3, 0, 0], chain_id="B", res_id=1, molecule_type_id=LIGAND),
    ]
    bonds = np.array(
        [
            (0, 1, BondType.SINGLE),  # intra-chain A (protein)
            (2, 3, BondType.SINGLE),  # intra-chain B (ligand)
            (1, 2, BondType.SINGLE),  # cross-chain covalent bond protein->ligand
        ]
    )
    atom_array = create_atomarray_with_bondlist(atoms, bonds)

    result = remove_covalent_nonprotein_chains(atom_array)

    # Ligand chain B should be removed, only protein chain A remains
    assert len(result) == 2
    assert set(result.chain_id) == {"A"}


def test_remove_covalent_nonprotein_chains_keeps_coordination_bonded():
    """Tests that coordination bonds don't trigger removal of non-protein chains."""
    # Protein chain A with coordination bond to metal ion chain B
    atoms = [
        Atom([0, 0, 0], chain_id="A", res_id=1, molecule_type_id=PROTEIN),
        Atom([2.3, 0, 0], chain_id="B", res_id=1, molecule_type_id=LIGAND),
    ]
    bonds = np.array(
        [
            (0, 1, BondType.COORDINATION),  # cross-chain coordination bond
        ]
    )
    atom_array = create_atomarray_with_bondlist(atoms, bonds)

    result = remove_covalent_nonprotein_chains(atom_array)

    # Both chains should remain since coordination bonds are excluded
    assert len(result) == 2
    assert set(result.chain_id) == {"A", "B"}


def test_remove_covalent_nonprotein_chains_keeps_protein_protein():
    """Tests that protein-protein cross-chain bonds don't trigger removal."""
    # Two protein chains with a disulfide-like bond
    atoms = [
        Atom([0, 0, 0], chain_id="A", res_id=1, molecule_type_id=PROTEIN),
        Atom([1, 0, 0], chain_id="A", res_id=1, molecule_type_id=PROTEIN),
        Atom([2, 0, 0], chain_id="B", res_id=1, molecule_type_id=PROTEIN),
        Atom([3, 0, 0], chain_id="B", res_id=1, molecule_type_id=PROTEIN),
    ]
    bonds = np.array(
        [
            (0, 1, BondType.SINGLE),  # intra-chain A
            (2, 3, BondType.SINGLE),  # intra-chain B
            (1, 2, BondType.SINGLE),  # cross-chain protein-protein bond
        ]
    )
    atom_array = create_atomarray_with_bondlist(atoms, bonds)

    result = remove_covalent_nonprotein_chains(atom_array)

    # Both protein chains should remain
    assert len(result) == 4
    assert set(result.chain_id) == {"A", "B"}


if __name__ == "__main__":
    pytest.main([__file__])
