# Usage:
# from warp_md.analysis.check_chirality import check_chirality
# out = check_chirality(traj, system, groups=[(":1@C", ":1@N", ":1@CA", ":1@CB")])

from __future__ import annotations

from typing import Iterable, Optional, Sequence, Tuple, Union

import numpy as np


MaskLike = Union[str, Sequence[int], np.ndarray, int]
GroupLike = Sequence[MaskLike]


def _read_all(traj, chunk_frames: Optional[int]):
    max_chunk = chunk_frames or 128
    coords_list = []
    chunk = traj.read_chunk(max_chunk)
    if chunk is None:
        return None
    while chunk is not None:
        coords_list.append(np.asarray(chunk["coords"], dtype=np.float64))
        chunk = traj.read_chunk(max_chunk)
    coords = np.concatenate(coords_list, axis=0) if coords_list else np.empty((0, 0, 3))
    return coords


def _indices_from_item(item: MaskLike) -> np.ndarray:
    if isinstance(item, np.ndarray):
        return np.asarray(item, dtype=np.int64).reshape(-1)
    if isinstance(item, (list, tuple)):
        return np.asarray([int(x) for x in item], dtype=np.int64)
    return np.asarray([int(item)], dtype=np.int64)


def _selection_from_item(system, item: MaskLike):
    if isinstance(item, str):
        return system.select(item)
    indices = _indices_from_item(item).tolist()
    return system.select_indices(indices)


def _center(coords: np.ndarray, idx: np.ndarray, masses: Optional[np.ndarray], mass_weighted: bool):
    sel = coords[:, idx, :]
    if sel.size == 0:
        return np.zeros((coords.shape[0], 3), dtype=np.float64)
    if mass_weighted and masses is not None and masses.size > 0:
        w = masses[idx]
        wsum = np.sum(w)
        if wsum <= 0.0:
            w = np.ones_like(w)
            wsum = np.sum(w)
        return (sel * w[None, :, None]).sum(axis=1) / wsum
    return sel.mean(axis=1)


def check_chirality(
    traj,
    system,
    groups: Iterable[GroupLike],
    mass_weighted: bool = False,
    chunk_frames: Optional[int] = None,
    device: str = "auto",
) -> np.ndarray:
    """Signed volume for chirality groups (A,B,C,D)."""
    use_plan = False
    try:
        from warp_md import CheckChiralityPlan  # type: ignore

        use_plan = getattr(CheckChiralityPlan, "__name__", "") != "_Missing"
    except Exception:
        CheckChiralityPlan = None  # type: ignore

    group_list = list(groups)
    if use_plan and CheckChiralityPlan is not None:
        sel_groups = []
        for group in group_list:
            if len(group) != 4:
                raise ValueError("each chirality group must have 4 selections")
            a = _selection_from_item(system, group[0])
            b = _selection_from_item(system, group[1])
            c = _selection_from_item(system, group[2])
            d = _selection_from_item(system, group[3])
            sel_groups.append((a, b, c, d))
        plan = CheckChiralityPlan(sel_groups, mass_weighted=mass_weighted)
        out = plan.run(traj, system, chunk_frames=chunk_frames, device=device)
        return np.asarray(out, dtype=np.float32)

    coords = _read_all(traj, chunk_frames)
    if coords is None:
        raise ValueError("trajectory has no frames")
    if coords.size == 0 or not group_list:
        return np.empty((0, 0), dtype=np.float32)

    masses = None
    if mass_weighted and hasattr(system, "atom_table"):
        atoms = system.atom_table()
        masses = np.asarray(atoms.get("mass", []), dtype=np.float64)

    n_frames = coords.shape[0]
    results = np.zeros((n_frames, len(group_list)), dtype=np.float32)
    for g_idx, group in enumerate(group_list):
        if len(group) != 4:
            raise ValueError("each chirality group must have 4 selections")
        idx_a = _indices_from_item(group[0])
        idx_b = _indices_from_item(group[1])
        idx_c = _indices_from_item(group[2])
        idx_d = _indices_from_item(group[3])
        pa = _center(coords, idx_a, masses, mass_weighted)
        pb = _center(coords, idx_b, masses, mass_weighted)
        pc = _center(coords, idx_c, masses, mass_weighted)
        pd = _center(coords, idx_d, masses, mass_weighted)
        ab = pb - pa
        ac = pc - pa
        ad = pd - pa
        cross = np.cross(ab, ac)
        vol = np.einsum("ij,ij->i", cross, ad)
        results[:, g_idx] = vol.astype(np.float32)
    return results


__all__ = ["check_chirality"]
