# Usage:
# from warp_md.analysis.symmrmsd import symmrmsd
# rms = symmrmsd(traj, system, mask="name CA", ref=0)

from __future__ import annotations

from typing import Optional, Sequence, Union

import numpy as np

from .rmsd import _kabsch_rmsd, _rmsd_raw, _read_all, _selection_indices


RefLike = Union[int, str]


def symmrmsd(
    traj,
    system,
    mask: str = "",
    ref: RefLike = 0,
    ref_mask: Optional[str] = None,
    fit: bool = True,
    remap: bool = False,
    mass: bool = False,
    frame_indices: Optional[Sequence[int]] = None,
    chunk_frames: Optional[int] = None,
) -> np.ndarray:
    """Symmetry-corrected RMSD (currently falls back to RMSD)."""
    if remap:
        raise ValueError("remap is not supported in this implementation")
    if mass:
        raise ValueError("mass-weighted symmrmsd not supported yet")

    coords, _box = _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)

    n_frames = coords.shape[0]
    if frame_indices is not None:
        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]
        n_frames = coords.shape[0]

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

    if ref_mask is None:
        ref_mask = mask
    ref_idx = _selection_indices(system, ref_mask)
    if ref_idx.size != idx.size:
        raise ValueError("mask and ref_mask selections must have same size")

    ref_coords = None
    if isinstance(ref, str):
        key = ref.strip().lower()
        if key in ("topology", "top", "topo"):
            if hasattr(system, "positions0"):
                ref_coords = system.positions0()
        if ref_coords is None and key in ("frame0", "first", "0", "topology", "top", "topo"):
            ref = 0

    if ref_coords is None:
        if isinstance(ref, int):
            ref_index = ref
            if ref_index < 0:
                ref_index = n_frames + ref_index
            if ref_index < 0 or ref_index >= n_frames:
                raise ValueError("ref index out of range")
            ref_coords = coords[ref_index]
        else:
            ref_coords = coords[0]
    else:
        ref_coords = np.asarray(ref_coords, dtype=np.float64)

    ref_sel = ref_coords[ref_idx]
    out = np.empty(n_frames, dtype=np.float64)
    for i in range(n_frames):
        cur = coords[i][idx]
        if fit:
            val = _kabsch_rmsd(cur, ref_sel)
        else:
            val = _rmsd_raw(cur, ref_sel)
        out[i] = val
    return out.astype(np.float32)


__all__ = ["symmrmsd"]
