# Usage:
# from warp_md.analysis.randomize_ions import randomize_ions
# traj_out = randomize_ions(traj, system, mask='@Na+', around=':1-16', by=5.0, overlap=3.0, seed=1)

from __future__ import annotations

from typing import Optional, Sequence

import numpy as np

from warp_md import RandomizeIonsPlan

from .trajectory import ArrayTrajectory


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")
        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 randomize_ions(
    traj,
    system,
    mask: str,
    around: Optional[str] = None,
    by: float = 0.0,
    overlap: float = 0.0,
    seed: int = 1,
    noimage: bool = False,
    frame_indices: Optional[Sequence[int]] = None,
    chunk_frames: Optional[int] = None,
    device: str = "auto",
):
    """Randomize ions by swapping with solvent residues."""
    if not mask:
        raise ValueError("mask is required")
    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)

    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)

    temp_traj = ArrayTrajectory(coords.astype(np.float32), box=box, time_ps=time)
    sel = system.select(mask)
    if len(sel.indices) == 0:
        return ArrayTrajectory(coords.astype(np.float32), box=box, time_ps=time)
    around_sel = system.select(around) if around else None

    plan = RandomizeIonsPlan(sel, seed, around_sel, float(by), float(overlap), bool(noimage))
    flat = plan.run(temp_traj, system, chunk_frames=None, device=device)
    arr = np.asarray(flat, dtype=np.float32)
    if arr.size == 0:
        return ArrayTrajectory(np.empty_like(coords, dtype=np.float32), box=box, time_ps=time)

    n_atoms = temp_traj.n_atoms()
    n_frames = arr.size // (n_atoms * 3)
    out_coords = arr.reshape((n_frames, n_atoms, 3))
    return ArrayTrajectory(out_coords, box=box, time_ps=time)


__all__ = ["randomize_ions"]
