import numpy as np

from warp_md.analysis.surf import molsurf, surf


class _DummySelection:
    def __init__(self, indices):
        self.indices = indices


class _DummySystem:
    def __init__(self, n_atoms):
        self._n_atoms = n_atoms

    def select(self, _mask):
        return _DummySelection(range(self._n_atoms))

    def select_indices(self, indices):
        return _DummySelection(indices)


class _DummyTraj:
    def __init__(self, coords):
        self._coords = coords
        self._used = False

    def read_chunk(self, _max_frames=128):
        if self._used:
            return None
        self._used = True
        return {"coords": self._coords}


def test_surf_bbox_area():
    coords = np.array(
        [
            [[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]],
            [[0.0, 0.0, 0.0], [2.0, 0.0, 0.0]],
        ],
        dtype=np.float64,
    )
    traj = _DummyTraj(coords)
    system = _DummySystem(coords.shape[1])

    out = surf(traj, system, mask=[0, 1])
    np.testing.assert_allclose(out, np.array([6.0, 0.0], dtype=np.float32), atol=1e-5)


def test_molsurf_alias():
    coords = np.array([[[0.0, 0.0, 0.0], [1.0, 2.0, 3.0]]], dtype=np.float64)
    traj = _DummyTraj(coords)
    system = _DummySystem(coords.shape[1])

    out = molsurf(traj, system, mask=[0, 1])
    dx, dy, dz = 1.0, 2.0, 3.0
    expected = 2.0 * (dx * dy + dy * dz + dx * dz)
    np.testing.assert_allclose(out, np.array([expected], dtype=np.float32), atol=1e-5)
