# SPDX-License-Identifier: LGPL-3.0-or-later
"""Tests for neighbor list functionality in torch-admp.

This module contains tests to verify the correctness of neighbor list
implementations, including comparisons with reference implementations
and consistency checks between different neighbor list methods.
"""

import unittest

import freud
import numpy as np
import torch

from torch_admp.nblist import (
    TorchNeighborList,
    dp_nblist,
    extend_input_and_build_neighbor_list,
    sort_pairs,
    vesin_nblist,
)
from torch_admp.utils import to_numpy_array

from . import SEED


class TestTorchNeighborList(unittest.TestCase):
    """Test class for TorchNeighborList functionality.

    Tests that the TorchNeighborList correctly identifies atom pairs
    within the cutoff distance, comparing against freud reference data.
    """

    def setUp(self):
        """Set up test data for TorchNeighborList testing.

        Creates reference neighbor list data using freud and initializes
        TorchNeighborList with test positions and box.
        """
        # Set random generators with SEED for reproducibility
        np_rng = np.random.default_rng(SEED)
        # torch_rng = torch.Generator().manual_seed(SEED)

        # reference data
        rcut = 4.0
        l_box = 10.0
        box = np.diag([l_box, l_box, l_box])
        positions = np_rng.random((20, 3)) * l_box

        fbox = freud.box.Box.from_matrix(box)
        aq = freud.locality.AABBQuery(fbox, positions)
        res = aq.query(positions, dict(r_max=rcut, exclude_ii=True))
        nblist = res.toNeighborList()
        nblist = np.vstack((nblist[:, 0], nblist[:, 1])).T
        nblist = nblist.astype(np.int32)
        msk = (nblist[:, 0] - nblist[:, 1]) < 0
        self.nblist_ref = nblist[msk]

        self.nblist = TorchNeighborList(rcut)
        self.positions = torch.tensor(positions)
        self.box = torch.tensor(box)

    def test_pairs(self):
        """Check that pairs are in the neighbor list.

        Verifies that all pairs identified by TorchNeighborList are
        present in the reference neighbor list from freud.
        """
        pairs = self.nblist(self.positions, self.box)
        pairs = to_numpy_array(pairs)
        mask = pairs[:, 0] < pairs[:, 1]
        assert len(pairs[mask]) == len(self.nblist_ref)
        for p in pairs[mask]:
            mask = (self.nblist_ref[:, 0] == p[0]) & (self.nblist_ref[:, 1] == p[1])
            self.assertTrue(mask.any())


class TestNBList(unittest.TestCase):
    """Test class for neighbor list consistency.

    Tests consistency between different neighbor list implementations
    to ensure they produce identical results.
    """

    def setUp(self) -> None:
        """Set up test data for neighbor list consistency testing.

        Creates random atomic positions in a periodic box and initializes
        parameters for neighbor list testing.
        """
        # Set random generators with SEED for reproducibility
        np_rng = np.random.default_rng(SEED)
        # torch_rng = torch.Generator().manual_seed(SEED)

        # reference data
        l_box = 10.0
        box = np.diag([l_box, l_box, l_box])
        positions = np_rng.random((100, 3)) * l_box

        self.positions = torch.tensor(positions)
        self.box = torch.tensor(box)

        # test: cutoff > l_box / 2!!!
        self.rcut = 6.0
        self.nnei = 150

    @unittest.skipIf(
        extend_input_and_build_neighbor_list is None, "deepmd.pt not installed"
    )
    def test_consistent(self):
        """Test consistency between different neighbor list implementations.

        Compares neighbor lists generated by dp_nblist and vesin_nblist
        to ensure they produce identical pairs and distances.
        """
        pairs_1, ds_1, _buffer_scales = dp_nblist(
            self.positions, self.box, self.nnei, self.rcut
        )
        pairs_2, ds_2, _buffer_scales = vesin_nblist(
            self.positions, self.box, self.rcut
        )
        torch.testing.assert_close(sort_pairs(pairs_1), sort_pairs(pairs_2))
        torch.testing.assert_close(torch.sort(ds_1)[0], torch.sort(ds_2)[0])
