import numpy as np

from warp_md.analysis.diffusion import tordiff


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))


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 _frame_for_angle(deg: float) -> np.ndarray:
    a = np.array([0.0, 1.0, 0.0])
    b = np.array([0.0, 0.0, 0.0])
    c = np.array([1.0, 0.0, 0.0])
    rad = np.deg2rad(deg)
    d = np.array([1.0, np.cos(rad), np.sin(rad)])
    return np.stack([a, b, c, d], axis=0)


def test_tordiff_bins():
    angles = [180.0, 0.0, 60.0, -60.0]
    coords = np.stack([_frame_for_angle(a) for a in angles], axis=0)
    traj = _DummyTraj(coords.astype(np.float64))
    system = _DummySystem(coords.shape[1])

    out = tordiff(traj, system, mask="all")
    expected = {
        "trans": [1.0, 0.0, 0.0, 0.0],
        "cis": [0.0, 1.0, 0.0, 0.0],
        "g_plus": [0.0, 0.0, 1.0, 0.0],
        "g_minus": [0.0, 0.0, 0.0, 1.0],
    }

    for key, values in expected.items():
        assert out[key].shape == (4,)
        np.testing.assert_allclose(out[key], np.asarray(values, dtype=np.float32), atol=1e-4)
    total = out["trans"] + out["cis"] + out["g_plus"] + out["g_minus"]
    np.testing.assert_allclose(total, np.ones_like(total), atol=1e-4)


def test_tordiff_output_files(tmp_path):
    angles = [180.0, 0.0]
    coords = np.stack([_frame_for_angle(a) for a in angles], axis=0)
    traj = _DummyTraj(coords.astype(np.float64))
    system = _DummySystem(coords.shape[1])

    out_path = tmp_path / "tordiff_out.txt"
    diff_path = tmp_path / "tordiff_diff.txt"
    out = tordiff(traj, system, mask="all", out=str(out_path), diffout=str(diff_path), time=2.0)
    assert out["time"].shape[0] == 2
    assert out_path.exists()
    assert diff_path.exists()
    data = np.loadtxt(out_path, comments="#")
    assert data.shape[0] == 2
