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

from __future__ import annotations

from typing import Optional, Sequence, Tuple, 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 _scale_box(box: Optional[np.ndarray], repeats: Tuple[int, int, int]) -> Optional[np.ndarray]:
    if box is None:
        return None
    if box.ndim == 2 and box.shape[1] == 3:
        return box * np.asarray(repeats, dtype=np.float64)[None, :]
    if box.ndim == 2 and box.shape[1] == 9:
        out = box.copy()
        out[:, 0:3] *= repeats[0]
        out[:, 3:6] *= repeats[1]
        out[:, 6:9] *= repeats[2]
        return out
    if box.ndim == 3 and box.shape[1:] == (3, 3):
        out = box.copy()
        out[:, :, 0] *= repeats[0]
        out[:, :, 1] *= repeats[1]
        out[:, :, 2] *= repeats[2]
        return out
    return None


def xtalsymm(
    traj,
    system,
    mask: MaskLike = "",
    repeats: Tuple[int, int, int] = (1, 1, 1),
    frame_indices: Optional[Sequence[int]] = None,
    chunk_frames: Optional[int] = None,
    device: str = "auto",
) -> ArrayTrajectory:
    """Replicate unit cell (selection only)."""
    reps = tuple(int(max(0, r)) for r in repeats)
    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)

    idx = _selection_indices(system, mask)
    if idx.size == 0 or reps == (0, 0, 0):
        return ArrayTrajectory(np.empty((coords.shape[0], 0, 3), dtype=np.float32), box=box, time_ps=time)

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

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

    if use_plan and XtalSymmPlan 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(idx.tolist())
        plan = XtalSymmPlan(sel, list(reps))
        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_frames = coords.shape[0]
        n_atoms = idx.size * reps[0] * reps[1] * reps[2]
        out = arr.reshape((n_frames, n_atoms, 3))
        new_box = _scale_box(box, reps)
        return ArrayTrajectory(out, box=new_box, time_ps=time)

    if box is None or (box.ndim == 2 and box.shape[1] not in (3, 9)) or box.ndim > 3:
        raise ValueError("xtalsymm fallback requires box lengths or 3x3 cell")
    out_frames = []
    for f in range(coords.shape[0]):
        sel = coords[f, idx, :]
        frame_coords = []
        if box.ndim == 2 and box.shape[1] == 3:
            lx, ly, lz = box[f]
            for ix in range(reps[0]):
                for iy in range(reps[1]):
                    for iz in range(reps[2]):
                        shift = np.array([ix * lx, iy * ly, iz * lz], dtype=np.float64)
                        frame_coords.append(sel + shift[None, :])
        else:
            if box.ndim == 3:
                m = box[f]
                ax, ay, az = m[0, 0], m[1, 0], m[2, 0]
                bx, by, bz = m[0, 1], m[1, 1], m[2, 1]
                cx, cy, cz = m[0, 2], m[1, 2], m[2, 2]
            else:
                m = box[f]
                ax, ay, az = m[0], m[1], m[2]
                bx, by, bz = m[3], m[4], m[5]
                cx, cy, cz = m[6], m[7], m[8]
            for ix in range(reps[0]):
                for iy in range(reps[1]):
                    for iz in range(reps[2]):
                        shift = np.array(
                            [ix * ax + iy * bx + iz * cx, ix * ay + iy * by + iz * cy, ix * az + iy * bz + iz * cz],
                            dtype=np.float64,
                        )
                        frame_coords.append(sel + shift[None, :])
        out_frames.append(np.concatenate(frame_coords, axis=0))
    out = np.stack(out_frames, axis=0) if out_frames else np.empty((0, 0, 3), dtype=np.float64)
    new_box = _scale_box(box, reps)
    return ArrayTrajectory(out.astype(np.float32), box=new_box, time_ps=time)


__all__ = ["xtalsymm"]
