# Usage:
# from warp_md.analysis.fiximagedbonds import fiximagedbonds
# new_traj = fiximagedbonds(traj, system, mask="all")

from __future__ import annotations

from typing import Optional, Sequence, Union

import numpy as np

from .trajectory import ArrayTrajectory


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 = []
    box_list = []
    time_list = []
    chunk = traj.read_chunk(max_chunk)
    if chunk is None:
        return None, None, None
    while chunk is not None:
        coords_list.append(np.asarray(chunk["coords"], dtype=np.float64))
        box = chunk.get("box")
        if box is not None:
            box_list.append(np.asarray(box, 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))
    box = np.concatenate(box_list, axis=0) if box_list else None
    time = np.concatenate(time_list, axis=0) if time_list else None
    return coords, box, time


def _apply_frame_indices(coords, box, time, frame_indices: Optional[Sequence[int]]):
    if frame_indices is None:
        return coords, box, time
    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 box is not None:
        box = box[select]
    if time is not None:
        time = time[select]
    return coords, box, time


def fiximagedbonds(
    traj,
    system,
    mask: MaskLike = "",
    frame_indices: Optional[Sequence[int]] = None,
    chunk_frames: Optional[int] = None,
    device: str = "auto",
) -> ArrayTrajectory:
    """Image selection into the primary unit cell."""
    coords, box, time = _read_all(traj, chunk_frames)
    if coords is None:
        raise ValueError("trajectory has no frames")
    if coords.size == 0:
        return ArrayTrajectory(coords.astype(np.float32), box=box, time_ps=time)
    if box is None:
        raise ValueError("fiximagedbonds requires box lengths")

    coords, box, time = _apply_frame_indices(coords, box, time, frame_indices)
    if coords.size == 0:
        return ArrayTrajectory(coords.astype(np.float32), box=box, time_ps=time)

    idx = _selection_indices(system, mask)
    if idx.size == 0:
        raise ValueError("selection resolved to empty set")

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

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

    if use_plan and FixImageBondsPlan is not None:
        sel = system.select(mask if isinstance(mask, str) else _all_resid_mask(system))
        if not isinstance(mask, str):
            sel = system.select_indices(idx.tolist())
        plan = FixImageBondsPlan(sel)
        temp = ArrayTrajectory(coords.astype(np.float32), box=box, time_ps=time)
        flat = plan.run(temp, system, chunk_frames=None, device=device)
        arr = np.asarray(flat, dtype=np.float32)
        n_atoms = system.n_atoms()
        n_frames = arr.size // (n_atoms * 3)
        out = arr.reshape((n_frames, n_atoms, 3))
        return ArrayTrajectory(out, box=box, time_ps=time)

    out = coords.copy()
    for f in range(out.shape[0]):
        lx, ly, lz = box[f]
        if lx <= 0 or ly <= 0 or lz <= 0:
            continue
        sel = out[f, idx, :]
        sel[:, 0] -= np.floor(sel[:, 0] / lx) * lx
        sel[:, 1] -= np.floor(sel[:, 1] / ly) * ly
        sel[:, 2] -= np.floor(sel[:, 2] / lz) * lz
        out[f, idx, :] = sel
    return ArrayTrajectory(out.astype(np.float32), box=box, time_ps=time)


__all__ = ["fiximagedbonds"]
