# Usage:
# from warp_md.analysis.dihedral_tools import rotate_dihedral, set_dihedral
# out = rotate_dihedral(traj, system, atoms=[0, 1, 2, 3], rotate_mask=[3], angle=30.0)

from __future__ import annotations

from typing import Optional, Sequence, Tuple, Union

import numpy as np

from .trajectory import ArrayTrajectory


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


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 _indices_from_mask(system, mask: MaskLike) -> np.ndarray:
    if isinstance(mask, np.ndarray):
        arr = np.asarray(mask, dtype=np.int64).reshape(-1)
        return arr
    if isinstance(mask, (list, tuple)):
        return np.asarray([int(x) for x in mask], dtype=np.int64)
    if not isinstance(mask, str):
        raise ValueError("mask must be string or sequence of indices")
    sel = system.select(mask)
    return np.asarray(list(sel.indices), dtype=np.int64)


def _selection_from_item(system, item: AtomLike):
    if isinstance(item, str):
        return system.select(item)
    try:
        idx = int(item)
    except Exception as exc:
        raise ValueError("atom entry must be int or selection string") from exc
    return system.select_indices([idx])


def _center(coords: np.ndarray, idx: np.ndarray, masses: Optional[np.ndarray], mass: bool) -> np.ndarray:
    sel = coords[:, idx, :]
    if sel.size == 0:
        raise ValueError("selection resolved to empty set")
    if mass and masses is not None and masses.size > 0:
        w = masses[idx]
        wsum = np.sum(w)
        if wsum <= 0.0:
            w = np.ones_like(w)
            wsum = np.sum(w)
        return (sel * w[None, :, None]).sum(axis=1) / wsum
    return sel.mean(axis=1)


def _apply_pbc(vec: np.ndarray, box: np.ndarray) -> np.ndarray:
    return vec - np.round(vec / box) * box


def _dihedral_from_points(p0, p1, p2, p3, degrees=True, range360=False) -> np.ndarray:
    b0 = p0 - p1
    b1 = p2 - p1
    b2 = p3 - p2
    b1_norm = np.linalg.norm(b1, axis=1)
    if np.any(b1_norm == 0.0):
        raise ValueError("dihedral axis has zero length")
    b1n = b1 / b1_norm[:, None]
    v = b0 - (np.einsum("ij,ij->i", b0, b1n))[:, None] * b1n
    w = b2 - (np.einsum("ij,ij->i", b2, b1n))[:, None] * b1n
    x = np.einsum("ij,ij->i", v, w)
    y = np.einsum("ij,ij->i", np.cross(b1n, v), w)
    ang = np.arctan2(y, x)
    if degrees:
        ang = np.degrees(ang)
        if range360:
            ang = ang % 360.0
    else:
        if range360:
            ang = ang % (2.0 * np.pi)
    return ang


def _angle_diff(target: float, current: np.ndarray, degrees: bool) -> np.ndarray:
    period = 360.0 if degrees else 2.0 * np.pi
    diff = (target - current + period / 2.0) % period - period / 2.0
    return diff


def _rotate_about_axis(coords: np.ndarray, idx: np.ndarray, pivot: np.ndarray, axis: np.ndarray, angle: np.ndarray):
    axis_norm = np.linalg.norm(axis, axis=1)
    if np.any(axis_norm == 0.0):
        raise ValueError("rotation axis has zero length")
    k = axis / axis_norm[:, None]
    v = coords[:, idx, :] - pivot[:, None, :]
    cosv = np.cos(angle)[:, None, None]
    sinv = np.sin(angle)[:, None, None]
    cross = np.cross(k[:, None, :], v)
    dot = np.sum(v * k[:, None, :], axis=2)[:, :, None]
    v_rot = v * cosv + cross * sinv + k[:, None, :] * dot * (1.0 - cosv)
    coords[:, idx, :] = pivot[:, None, :] + v_rot


def rotate_dihedral(
    traj,
    system,
    atoms: Sequence[AtomLike],
    rotate_mask: Optional[MaskLike] = None,
    angle: float = 0.0,
    mass: bool = False,
    degrees: bool = True,
    chunk_frames: Optional[int] = None,
    device: str = "auto",
):
    """Rotate a selection around the B-C dihedral axis."""
    if len(atoms) != 4:
        raise ValueError("atoms must contain 4 items")

    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)

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

        use_plan = hasattr(system, "select_indices")
    except Exception:
        RotateDihedralPlan = None  # type: ignore

    if use_plan and RotateDihedralPlan is not None:
        sel_a = _selection_from_item(system, atoms[0])
        sel_b = _selection_from_item(system, atoms[1])
        sel_c = _selection_from_item(system, atoms[2])
        sel_d = _selection_from_item(system, atoms[3])
        if rotate_mask is None:
            rotate_sel = sel_d
        else:
            if isinstance(rotate_mask, str):
                rotate_sel = system.select(rotate_mask)
            else:
                rotate_sel = system.select_indices(list(np.asarray(rotate_mask, dtype=np.int64)))
        plan = RotateDihedralPlan(sel_a, sel_b, sel_c, sel_d, rotate_sel, float(angle), mass, degrees)
        temp_traj = ArrayTrajectory(coords.astype(np.float32), box=box, time_ps=time)
        flat = plan.run(temp_traj, 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)

    masses = None
    if mass:
        atoms_tbl = system.atom_table()
        masses = np.asarray(atoms_tbl.get("mass", []), dtype=np.float64)

    idx_a = _indices_from_mask(system, atoms[0]) if isinstance(atoms[0], str) else np.asarray([int(atoms[0])])
    idx_b = _indices_from_mask(system, atoms[1]) if isinstance(atoms[1], str) else np.asarray([int(atoms[1])])
    idx_c = _indices_from_mask(system, atoms[2]) if isinstance(atoms[2], str) else np.asarray([int(atoms[2])])
    idx_d = _indices_from_mask(system, atoms[3]) if isinstance(atoms[3], str) else np.asarray([int(atoms[3])])
    if rotate_mask is None:
        rot_idx = idx_d
    else:
        rot_idx = _indices_from_mask(system, rotate_mask)

    pivot = _center(coords, idx_b, masses, mass)
    axis = _center(coords, idx_c, masses, mass) - pivot
    ang = float(angle)
    if degrees:
        ang = np.deg2rad(ang)
    angles = np.full((coords.shape[0],), ang, dtype=np.float64)
    out = coords.copy()
    _rotate_about_axis(out, rot_idx, pivot, axis, angles)
    return ArrayTrajectory(out.astype(np.float32), box=box, time_ps=time)


def set_dihedral(
    traj,
    system,
    atoms: Sequence[AtomLike],
    rotate_mask: Optional[MaskLike] = None,
    target: float = 0.0,
    mass: bool = False,
    pbc: str = "none",
    degrees: bool = True,
    range360: bool = False,
    chunk_frames: Optional[int] = None,
    device: str = "auto",
):
    """Rotate selection to set dihedral to a target angle."""
    if len(atoms) != 4:
        raise ValueError("atoms must contain 4 items")

    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)

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

        use_plan = hasattr(system, "select_indices")
    except Exception:
        SetDihedralPlan = None  # type: ignore

    if use_plan and SetDihedralPlan is not None:
        sel_a = _selection_from_item(system, atoms[0])
        sel_b = _selection_from_item(system, atoms[1])
        sel_c = _selection_from_item(system, atoms[2])
        sel_d = _selection_from_item(system, atoms[3])
        if rotate_mask is None:
            rotate_sel = sel_d
        else:
            if isinstance(rotate_mask, str):
                rotate_sel = system.select(rotate_mask)
            else:
                rotate_sel = system.select_indices(list(np.asarray(rotate_mask, dtype=np.int64)))
        plan = SetDihedralPlan(
            sel_a,
            sel_b,
            sel_c,
            sel_d,
            rotate_sel,
            float(target),
            mass,
            pbc,
            degrees,
            range360,
        )
        temp_traj = ArrayTrajectory(coords.astype(np.float32), box=box, time_ps=time)
        flat = plan.run(temp_traj, 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)

    masses = None
    if mass:
        atoms_tbl = system.atom_table()
        masses = np.asarray(atoms_tbl.get("mass", []), dtype=np.float64)

    idx_a = _indices_from_mask(system, atoms[0]) if isinstance(atoms[0], str) else np.asarray([int(atoms[0])])
    idx_b = _indices_from_mask(system, atoms[1]) if isinstance(atoms[1], str) else np.asarray([int(atoms[1])])
    idx_c = _indices_from_mask(system, atoms[2]) if isinstance(atoms[2], str) else np.asarray([int(atoms[2])])
    idx_d = _indices_from_mask(system, atoms[3]) if isinstance(atoms[3], str) else np.asarray([int(atoms[3])])
    if rotate_mask is None:
        rot_idx = idx_d
    else:
        rot_idx = _indices_from_mask(system, rotate_mask)

    pivot = _center(coords, idx_b, masses, mass)
    axis = _center(coords, idx_c, masses, mass) - pivot
    a = _center(coords, idx_a, masses, mass)
    d = _center(coords, idx_d, masses, mass)
    if pbc.lower() == "orthorhombic":
        if box is None:
            raise ValueError("orthorhombic PBC requires box lengths")
        b0 = a - pivot
        b1 = _center(coords, idx_c, masses, mass) - pivot
        b2 = d - _center(coords, idx_c, masses, mass)
        b0 = _apply_pbc(b0, box)
        b1 = _apply_pbc(b1, box)
        b2 = _apply_pbc(b2, box)
        a = pivot + b0
        d = _center(coords, idx_c, masses, mass) + b2

    current = _dihedral_from_points(a, pivot, _center(coords, idx_c, masses, mass), d, degrees=degrees, range360=range360)
    diff = _angle_diff(float(target), current, degrees)
    if degrees:
        diff = np.deg2rad(diff)
    out = coords.copy()
    _rotate_about_axis(out, rot_idx, pivot, axis, diff)
    return ArrayTrajectory(out.astype(np.float32), box=box, time_ps=time)


__all__ = ["rotate_dihedral", "set_dihedral"]
