# =======================
# Python internal package
# =======================
# P
import pathlib

# W
import warnings

# ================
# External package
# ================
# N
import numpy

# P
import pytest

# ==============
# Module package
# ==============
# P
from src.strange.parser.molecule import (
    Molecule,
    MoleculeParser,
    planarity_computation,
)


@pytest.mark.parametrize(
    "file_extension",
    [
        "mol",
        "mol2",
        "pdb",
        "pdbqt",
        "sdf",
        "xyz",
    ],
)
def test_file_readable(file_extension: str) -> None:
    """Test whether molecular files of various formats are readable.

    Parameters
    ----------
    file_extension : `str`

        The file extension of the molecular file to be tested.
    """
    path: pathlib.Path = pathlib.Path(
        f"tests/data/aspirin/molecule.{file_extension}"
    )

    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=UserWarning)
        assert isinstance(MoleculeParser(path), Molecule)


def test_add_protons() -> None:
    """Test proton addition during molecule parsing."""
    path: pathlib.Path = pathlib.Path(
        "tests/data/aspirin/molecule_deprotonated.pdb"
    )

    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=UserWarning)
        assert isinstance(MoleculeParser(path, add_hydrogen=True), Molecule)


@pytest.fixture
def __molecule() -> Molecule:
    """Provide a parsed aspirin molecule for testing.

    Returns
    -------
    `Molecule`
        A parsed molecular structure used across multiple tests.
    """
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=UserWarning)

        return MoleculeParser(pathlib.Path("tests/data/aspirin/molecule.pdb"))


@pytest.mark.parametrize(
    "coordinate, index",
    [
        (numpy.array([1.233, 0.554, 0.779]), 0),
        (numpy.array([-0.695, -2.715, -0.750]), 1),
        (numpy.array([0.796, -2.184, 0.868]), 2),
        (numpy.array([1.781, 0.811, -1.482]), 3),
        (numpy.array([-0.086, 0.609, 0.440]), 4),
        (numpy.array([-0.793, -0.552, 0.124]), 5),
        (numpy.array([-0.729, 1.846, 0.413]), 6),
        (numpy.array([-2.143, -0.474, -0.218]), 7),
        (numpy.array([-2.079, 1.924, 0.071]), 8),
        (numpy.array([-2.786, 0.764, -0.245]), 9),
        (numpy.array([-0.141, -1.854, 0.148]), 10),
        (numpy.array([2.109, 0.672, -0.311]), 11),
        (numpy.array([3.530, 0.600, 0.163]), 12),
        (numpy.array([-0.185, 2.754, 0.659]), 13),
        (numpy.array([-2.725, -1.360, -0.456]), 14),
        (numpy.array([-2.580, 2.887, 0.051]), 15),
        (numpy.array([-3.837, 0.824, -0.509]), 16),
        (numpy.array([3.729, 1.418, 0.859]), 17),
        (numpy.array([4.205, 0.697, -0.692]), 18),
        (numpy.array([3.711, -0.366, 0.643]), 19),
        (numpy.array([-0.255, -3.592, -0.734]), 20),
    ],
)
def test_molecule_get_index(
    coordinate: numpy.ndarray, index: int, __molecule: Molecule
) -> None:
    """Test atom index retrieval from spatial coordinates.

    Parameters
    ----------
    coordinate : `numpy.ndarray`
        The 3D coordinate used to query the molecule.

    index : `int`
        The expected atom index corresponding to the coordinate.

    __molecule : `Molecule`
        The molecule instance under test.
    """
    assert __molecule.get_index(coordinate) == index


@pytest.mark.parametrize(
    "index",
    range(21),
)
def test_molecule_get_data(index: int, __molecule: Molecule) -> None:
    """Test retrieval of atom metadata from a molecule.

    Parameters
    ----------
    index : `int`
        The atom index to retrieve data for.

    __molecule : `Molecule`
        The molecule instance under test.
    """
    PATH: str = "tests/data/aspirin/molecule.pdb"
    data: dict[str, str | int] = {"chain": "SYSTEM"}

    with open(PATH, "r", encoding="utf-8") as file:
        for line_index, line in enumerate(file):
            if line_index != index:
                continue

            data["atom_name"] = line[12:16].strip()
            data["atom_type"] = line[76:78].strip()
            data["atom_id"] = int(line[6:11].strip())
            data["residu_name"] = line[17:20].strip()
            data["residu_id"] = int(line[22:26].strip())

            break

    assert __molecule.get_data(index) == data


@pytest.mark.parametrize(
    "coordinate, expected_overlap",
    [
        (
            numpy.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [-1.0, 0.0, 0.0]]),
            numpy.array([0.0, 0.0, 0.0]),
        ),
        (
            numpy.array(
                [[6.0, 0.0, 0.5], [7.0, -8.0, -0.5], [1.8, 65.0, 0.5]]
            ),
            numpy.array([-1.16246265e-17, -3.85833298e-18, 2.20857500e-17]),
        ),
    ],
)
def test_planarity_computation(
    coordinate: numpy.ndarray, expected_overlap: numpy.ndarray
) -> None:
    """Test computation of molecular planarity overlap.

    Parameters
    ----------
    coordinate : `numpy.ndarray`
        A set of 3D coordinates representing atom positions.

    expected_overlap : `numpy.ndarray`
        The expected planarity overlap vector.
    """
    overlap = planarity_computation(coordinate)

    assert overlap.shape == (3,)
    assert overlap == pytest.approx(expected_overlap, abs=1e-5)
