# Usage:
# from warp_md.analysis.rotdif import rotdif
# time, data = rotdif(traj, system, mask=":1-10", orientation=[1, 2], group_by="resid")

from __future__ import annotations

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

import numpy as np


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


def _all_resid_mask(system) -> str:
    if not hasattr(system, "atom_table"):
        return "all"
    atoms = system.atom_table()
    resids = atoms.get("resid", [])
    if not resids:
        return "resid 0:0"
    return f"resid {min(resids)}:{max(resids)}"


def _selection_indices(system, mask: MaskLike) -> np.ndarray:
    if isinstance(mask, np.ndarray):
        return np.asarray(mask, dtype=np.int64).reshape(-1)
    if isinstance(mask, (list, tuple)):
        return np.asarray([int(x) for x in mask], dtype=np.int64)
    if mask in ("", "*", "all", None):
        sel = system.select(_all_resid_mask(system))
    else:
        sel = system.select(mask)
    return np.asarray(list(sel.indices), dtype=np.int64)


def _read_all(traj, chunk_frames: Optional[int]):
    max_chunk = chunk_frames or 128
    coords_list = []
    time_list = []
    chunk = traj.read_chunk(max_chunk)
    if chunk is None:
        return None, None
    while chunk is not None:
        coords_list.append(np.asarray(chunk["coords"], dtype=np.float64))
        time = chunk.get("time") or chunk.get("time_ps")
        if time is not None:
            time_list.append(np.asarray(time, 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))
    times = np.concatenate(time_list, axis=0) if time_list else None
    return coords, times


def _apply_frame_indices(coords, times, frame_indices: Optional[Sequence[int]]):
    if frame_indices is None:
        return coords, times
    n_frames = coords.shape[0]
    select = []
    for i in frame_indices:
        j = int(i)
        if j < 0:
            j = n_frames + j
        if 0 <= j < n_frames:
            select.append(j)
    coords = coords[select]
    if times is not None:
        times = times[select]
    return coords, times


def _group_indices(system, indices: np.ndarray, group_by: str) -> Sequence[np.ndarray]:
    if group_by in ("all", "system"):
        return [indices.astype(np.int64)]
    if group_by == "resid" and hasattr(system, "atom_table"):
        atoms = system.atom_table()
        resids = atoms.get("resid", [])
        if not resids:
            return [indices.astype(np.int64)]
        groups = {}
        for idx in indices:
            resid = resids[int(idx)]
            groups.setdefault(resid, []).append(int(idx))
        return [np.asarray(groups[k], dtype=np.int64) for k in sorted(groups)]
    raise ValueError("fallback rotdif supports group_by='all' or 'resid'")


def _unit(v: np.ndarray) -> np.ndarray:
    norm = np.linalg.norm(v, axis=1, keepdims=True)
    norm = np.where(norm == 0.0, 1.0, norm)
    return v / norm


def _orientation_vectors(coords: np.ndarray, groups: Sequence[np.ndarray], orientation: Sequence[int]) -> np.ndarray:
    idxs = [int(i) for i in orientation]
    if len(idxs) not in (2, 3):
        raise ValueError("orientation must have 2 (vector) or 3 (plane) indices")
    max_idx = max(idxs)
    n_frames = coords.shape[0]
    n_groups = len(groups)
    out = np.zeros((n_frames, n_groups, 3), dtype=np.float64)
    for g, atoms in enumerate(groups):
        if atoms.size < max_idx:
            raise ValueError("orientation index out of range for group")
        a = atoms[idxs[0] - 1]
        b = atoms[idxs[1] - 1]
        if len(idxs) == 2:
            v = coords[:, b, :] - coords[:, a, :]
            out[:, g, :] = _unit(v)
        else:
            c = atoms[idxs[2] - 1]
            v1 = coords[:, b, :] - coords[:, a, :]
            v2 = coords[:, c, :] - coords[:, a, :]
            cross = np.cross(v1, v2)
            out[:, g, :] = _unit(cross)
    return out


def _rotacf_series(vectors: np.ndarray, max_lag: Optional[int], p2_legendre: bool) -> Tuple[np.ndarray, np.ndarray]:
    n_frames = vectors.shape[0]
    if n_frames == 0:
        return np.empty((0,), dtype=np.float32), np.empty((0, 2), dtype=np.float32)
    lag_max = max_lag if max_lag is not None else n_frames - 1
    lag_max = max(0, min(int(lag_max), n_frames - 1))
    rows = lag_max + 1
    out = np.zeros((rows, 2), dtype=np.float32)
    for lag in range(rows):
        v0 = vectors[: n_frames - lag]
        v1 = vectors[lag:]
        dot = np.sum(v0 * v1, axis=2)
        p1 = np.mean(dot)
        p2 = np.mean(dot * dot)
        if p2_legendre:
            p2 = 1.5 * p2 - 0.5
        out[lag, 0] = float(p1)
        out[lag, 1] = float(p2)
    time = np.arange(rows, dtype=np.float32)
    return time, out


def rotdif(
    traj,
    system,
    mask: MaskLike = "",
    orientation: Optional[Sequence[int]] = None,
    group_by: str = "resid",
    p2_legendre: bool = True,
    length_scale: Optional[float] = None,
    frame_decimation: Optional[Tuple[int, int]] = None,
    dt_decimation: Optional[Tuple[int, int, int, int]] = None,
    time_binning: Optional[Tuple[float, float]] = None,
    group_types: Optional[Sequence[int]] = None,
    lag_mode: Optional[str] = None,
    max_lag: Optional[int] = None,
    memory_budget_bytes: Optional[int] = None,
    multi_tau_m: Optional[int] = None,
    multi_tau_levels: Optional[int] = None,
    frame_indices: Optional[Sequence[int]] = None,
    chunk_frames: Optional[int] = None,
    device: str = "auto",
    dt: float = 1.0,
) -> Tuple[np.ndarray, np.ndarray]:
    """Rotational autocorrelation (RotAcfPlan wrapper)."""
    if orientation is None:
        raise ValueError("orientation indices required")

    use_plan = False
    try:
        from warp_md import RotAcfPlan  # type: ignore

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

    if use_plan and RotAcfPlan is not None:
        if isinstance(mask, str):
            sel = system.select(mask if mask not in ("", "*", "all", None) else _all_resid_mask(system))
        else:
            sel = system.select_indices(np.asarray(mask, dtype=np.int64).reshape(-1).tolist())
        plan = RotAcfPlan(
            sel,
            group_by=group_by,
            orientation=list(orientation),
            p2_legendre=p2_legendre,
            length_scale=length_scale,
            frame_decimation=frame_decimation,
            dt_decimation=dt_decimation,
            time_binning=time_binning,
            group_types=list(group_types) if group_types is not None else None,
            lag_mode=lag_mode,
            max_lag=max_lag,
            memory_budget_bytes=memory_budget_bytes,
            multi_tau_m=multi_tau_m,
            multi_tau_levels=multi_tau_levels,
        )
        time, data = plan.run(traj, system, chunk_frames=chunk_frames, device=device)
        return np.asarray(time, dtype=np.float32), np.asarray(data, dtype=np.float32)

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

    coords, times = _apply_frame_indices(coords, times, frame_indices)
    if coords.size == 0:
        return np.empty((0,), dtype=np.float32), np.empty((0, 2), dtype=np.float32)

    idx = _selection_indices(system, mask)
    if idx.size == 0:
        return np.empty((0,), dtype=np.float32), np.empty((0, 2), dtype=np.float32)

    groups = _group_indices(system, idx, group_by)
    vectors = _orientation_vectors(coords, groups, orientation)
    time, data = _rotacf_series(vectors, max_lag, p2_legendre)

    if times is not None and times.size >= 2:
        dt0 = float(np.median(np.diff(times)))
        time = time * dt0
    else:
        time = time * float(dt)
    return time.astype(np.float32), data.astype(np.float32)


__all__ = ["rotdif"]
