import sys
from pathlib import Path

import numpy as np
import pytest

ROOT = Path(__file__).resolve().parents[3]
sys.path.insert(0, str(ROOT / "python"))


def _require_openmm():
    try:
        from openmm import openmm as mm
        from openmm.app import Topology, element
    except Exception:  # pragma: no cover
        pytest.skip("openmm not available")
    return mm, Topology, element


class _FakeTraj:
    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 _build_system_and_topology():
    mm, Topology, element = _require_openmm()
    top = Topology()
    chain = top.addChain()
    res_sol = top.addResidue("LIG", chain)
    top.addAtom("C1", element.carbon, res_sol)
    res_wat = top.addResidue("HOH", chain)
    top.addAtom("O", element.oxygen, res_wat)
    top.addAtom("H1", element.hydrogen, res_wat)
    top.addAtom("H2", element.hydrogen, res_wat)

    system = mm.System()
    for _ in range(4):
        system.addParticle(1.0)
    nb = mm.NonbondedForce()
    for _ in range(4):
        nb.addParticle(0.0, 0.3, 0.0)
    nb.setNonbondedMethod(mm.NonbondedForce.NoCutoff)
    system.addForce(nb)
    return system, top


def test_gist_frame_indices():
    from warp_md.analysis.gist import GistConfig, gist

    system, top = _build_system_and_topology()
    coords = np.array(
        [
            [[0.0, 0.0, 0.0], [0.1, 0.0, 0.0], [0.1, 0.05, 0.0], [0.1, -0.05, 0.0]],
            [[0.0, 0.0, 0.0], [0.2, 0.0, 0.0], [0.2, 0.05, 0.0], [0.2, -0.05, 0.0]],
        ],
        dtype=np.float64,
    )
    traj = _FakeTraj(coords)
    config = GistConfig(length_scale=1.0, frame_indices=[1], grid_spacing=0.2, padding=0.4)
    out = gist(traj, None, system, top, config=config)
    assert out.n_frames == 1
    assert int(out.counts.sum()) == 1
