import numpy as np

from warp_md.analysis.xtalsymm import xtalsymm


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


class _DummySystem:
    def __init__(self, n_atoms):
        self._n_atoms = n_atoms
        self._atoms = {"resid": list(range(n_atoms))}

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

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

    def atom_table(self):
        return self._atoms

    def n_atoms(self):
        return self._n_atoms


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

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


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

    out = xtalsymm(traj, system, mask="all", repeats=(2, 1, 1))
    chunk = out.read_chunk()
    out_coords = chunk["coords"][0]
    assert out_coords.shape == (2, 3)
    np.testing.assert_allclose(out_coords[0], np.array([0.0, 0.0, 0.0], dtype=np.float32))
    np.testing.assert_allclose(out_coords[1], np.array([1.0, 0.0, 0.0], dtype=np.float32))
