import tempfile

import numpy as np

from warp_md.analysis.nmr import ired_vector_and_matrix, jcoupling, nh_order_parameters


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


class _DummySystem:
    def __init__(self, n_atoms):
        self._atoms = {
            "name": ["N", "H"] * (n_atoms // 2),
            "resname": ["ALA"] * n_atoms,
            "resid": [1] * n_atoms,
            "chain_id": [0] * n_atoms,
        }

    def atom_table(self):
        return self._atoms

    def select(self, _mask):
        return _DummySelection(range(len(self._atoms["name"])))


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_ired_frame_indices():
    coords = np.array(
        [
            [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]],
            [[0.0, 0.0, 0.0], [0.0, 1.0, 0.0]],
            [[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]],
        ],
        dtype=np.float64,
    )
    traj = _DummyTraj(coords)
    system = _DummySystem(coords.shape[1])
    vecs, mat = ired_vector_and_matrix(
        traj,
        system,
        selection="all",
        frame_indices=[1],
    )
    assert vecs.shape[0] == 1
    assert mat.shape[0] == 3


def test_nh_order_parameters_shape():
    coords = np.array(
        [
            [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]],
            [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]],
        ],
        dtype=np.float64,
    )
    traj = _DummyTraj(coords)
    system = _DummySystem(coords.shape[1])
    s2 = nh_order_parameters(traj, system, selection="all", frame_indices=[0, 1])
    assert s2.shape == (1,)


def test_jcoupling_kfile():
    coords = np.array(
        [
            [
                [1.0, 0.0, 0.0],
                [0.0, 0.0, 0.0],
                [0.0, 1.0, 0.0],
                [1.0, 1.0, 0.0],
            ]
        ],
        dtype=np.float64,
    )
    traj = _DummyTraj(coords)
    system = _DummySystem(coords.shape[1])
    with tempfile.NamedTemporaryFile("w+", delete=False) as handle:
        handle.write("1.0 2.0 3.0\n")
        kfile = handle.name
    values = jcoupling(
        traj,
        system,
        dihedral_indices=[(0, 1, 2, 3)],
        kfile=kfile,
    )
    assert values.shape == (1, 1)
    assert np.allclose(values[0, 0], 6.0, atol=1e-5)


def test_jcoupling_frame_indices():
    coords = np.array(
        [
            [
                [1.0, 0.0, 0.0],
                [0.0, 0.0, 0.0],
                [0.0, 1.0, 0.0],
                [1.0, 1.0, 0.0],
            ],
            [
                [1.0, 0.0, 0.0],
                [0.0, 0.0, 0.0],
                [0.0, 1.0, 0.0],
                [2.0, 1.0, 0.0],
            ],
        ],
        dtype=np.float64,
    )
    traj = _DummyTraj(coords)
    system = _DummySystem(coords.shape[1])
    values = jcoupling(
        traj,
        system,
        dihedral_indices=[(0, 1, 2, 3)],
        frame_indices=[1],
    )
    assert values.shape == (1, 1)
