# Usage:
# from warp_md.analysis.nmr import nh_order_parameters, jcoupling
# s2 = nh_order_parameters(traj, system, nh_pairs, order=2)
# j = jcoupling(traj, system, dih_indices, karplus=(6.4, -1.4, 1.9))

from __future__ import annotations

from typing import Iterable, Optional, Sequence, Tuple, Union

import numpy as np

PairLike = Union[Tuple[int, int], Sequence[int]]
QuadLike = Union[Tuple[int, int, int, int], Sequence[int]]


def _axis_norm(vec: np.ndarray) -> np.ndarray:
    norms = np.linalg.norm(vec, axis=-1, keepdims=True)
    norms = np.where(norms == 0.0, 1.0, norms)
    return vec / norms


def _pairs_from_selection(system, selection: str) -> Sequence[Tuple[int, int]]:
    sel = system.select(selection)
    idx = list(sel.indices)
    if len(idx) < 2:
        raise ValueError("selection must contain at least two atoms")
    return [(int(idx[i]), int(idx[i + 1])) for i in range(len(idx) - 1)]


def _resolve_pairs(system, vector_pairs, selection: Optional[str]):
    if vector_pairs is None:
        if not selection:
            raise ValueError("selection is required when vector_pairs is None")
        return _pairs_from_selection(system, selection)
    if isinstance(vector_pairs, str):
        if vector_pairs.lower() != "sequential":
            raise ValueError("vector_pairs string must be 'sequential'")
        if not selection:
            raise ValueError("selection is required for sequential pairs")
        return _pairs_from_selection(system, selection)
    pairs = []
    for pair in vector_pairs:
        if len(pair) != 2:
            raise ValueError("each vector pair must have length 2")
        a, b = pair
        if isinstance(a, str) and isinstance(b, str):
            sel_a = system.select(a)
            sel_b = system.select(b)
            if len(sel_a.indices) != len(sel_b.indices):
                raise ValueError("selection pairs must match in length")
            pairs.extend(zip((int(i) for i in sel_a.indices), (int(i) for i in sel_b.indices)))
        else:
            pairs.append((int(a), int(b)))
    return pairs


def ired_vector_and_matrix(
    traj,
    system,
    vector_pairs: Optional[Sequence[PairLike]] = None,
    selection: Optional[str] = None,
    order: int = 2,
    length_scale: float = 0.1,
    pbc: str = "none",
    chunk_frames: Optional[int] = None,
    frame_indices: Optional[Sequence[int]] = None,
):
    """Return unit vectors and correlation matrix for IRED-like workflow."""
    pbc_mode = pbc.lower()
    if pbc_mode not in ("none", "orthorhombic"):
        raise ValueError("pbc must be 'none' or 'orthorhombic'")
    pairs = _resolve_pairs(system, vector_pairs, selection)
    idx_a = np.array([p[0] for p in pairs], dtype=np.int64)
    idx_b = np.array([p[1] for p in pairs], dtype=np.int64)
    n_pairs = len(pairs)

    vecs = []
    max_chunk = chunk_frames or 128
    chunk = traj.read_chunk(max_chunk)
    if chunk is None:
        raise ValueError("trajectory has no frames")
    if frame_indices is not None:
        frame_set = {int(i) for i in frame_indices if int(i) >= 0}
    else:
        frame_set = None
    global_frame = 0

    while chunk is not None:
        coords = np.asarray(chunk["coords"], dtype=np.float64) * length_scale
        box = chunk.get("box")
        if box is not None:
            box = np.asarray(box, dtype=np.float64) * length_scale
        frames = coords.shape[0]
        for f in range(frames):
            if frame_set is not None and global_frame not in frame_set:
                global_frame += 1
                continue
            v = coords[f, idx_b] - coords[f, idx_a]
            if pbc_mode == "orthorhombic":
                if box is None:
                    raise ValueError("pbc='orthorhombic' requires box in trajectory")
                b = box[f]
                if np.any(b == 0.0):
                    raise ValueError("box lengths must be non-zero for pbc")
                v = v - np.round(v / b) * b
            vecs.append(_axis_norm(v))
            global_frame += 1
        chunk = traj.read_chunk(max_chunk)
    vecs = np.asarray(vecs, dtype=np.float64)  # (n_frames, n_pairs, 3)
    if vecs.size == 0:
        return vecs, np.empty((0, 0), dtype=np.float64)
    # build correlation matrix
    n_frames = vecs.shape[0]
    flat = vecs.reshape(n_frames, n_pairs * 3)
    mat = (flat.T @ flat) / float(n_frames)
    if order == 2:
        # normalize by first element (approx)
        if mat.shape[0] > 0 and mat[0, 0] != 0.0:
            mat = mat / mat[0, 0]
    return vecs, mat


def nh_order_parameters(
    traj,
    system,
    vector_pairs: Optional[Sequence[PairLike]] = None,
    selection: Optional[str] = None,
    order: int = 2,
    tstep: float = 1.0,
    tcorr: float = 10000.0,
    length_scale: float = 0.1,
    pbc: str = "none",
    chunk_frames: Optional[int] = None,
    frame_indices: Optional[Sequence[int]] = None,
):
    """Compute approximate N-H order parameters (S2)."""
    vecs, _ = ired_vector_and_matrix(
        traj,
        system,
        vector_pairs=vector_pairs,
        selection=selection,
        order=order,
        length_scale=length_scale,
        pbc=pbc,
        chunk_frames=chunk_frames,
        frame_indices=frame_indices,
    )
    if vecs.size == 0:
        return np.empty(0, dtype=np.float32)
    # Order tensor approach: S2 = 1.5*Tr(M^2) - 0.5 where M = <u_i u_j>
    n_frames = vecs.shape[0]
    m = np.einsum("fpi,fpj->pij", vecs, vecs) / float(n_frames)
    tr_m2 = np.einsum("pij,pij->p", m, m)
    s2 = 1.5 * tr_m2 - 0.5
    return s2.astype(np.float32)


def jcoupling(
    traj,
    system,
    dihedral_indices: Sequence[QuadLike],
    karplus: Tuple[float, float, float] = (6.4, -1.4, 1.9),
    kfile: Optional[str] = None,
    phase_deg: float = 0.0,
    length_scale: float = 0.1,
    pbc: str = "none",
    chunk_frames: Optional[int] = None,
    frame_indices: Optional[Sequence[int]] = None,
    device: str = "auto",
):
    """Compute scalar J-coupling via a Karplus-like equation.

    J = A*cos^2(theta+phase) + B*cos(theta+phase) + C
    """
    if kfile is not None:
        with open(kfile, "r", encoding="utf-8") as handle:
            found = None
            for line in handle:
                line = line.strip()
                if not line or line.startswith("#") or line.startswith(";"):
                    continue
                parts = line.split()
                nums = []
                for part in parts:
                    try:
                        nums.append(float(part))
                    except ValueError:
                        break
                if len(nums) >= 3:
                    found = (nums[0], nums[1], nums[2])
                    break
            if found is None:
                raise ValueError("kfile does not contain 3 numeric karplus parameters")
        a, b, c = found
    else:
        a, b, c = karplus
    phase = np.deg2rad(phase_deg)
    resolved = []
    for quad in dihedral_indices:
        if len(quad) != 4:
            raise ValueError("each dihedral entry must have length 4")
        if all(isinstance(v, str) for v in quad):
            sels = [system.select(v) for v in quad]
            if any(len(s.indices) != 1 for s in sels):
                raise ValueError("each selection in dihedral entry must map to one atom")
            resolved.append(tuple(int(s.indices[0]) for s in sels))
        else:
            resolved.append(tuple(int(v) for v in quad))
    idx = np.array(resolved, dtype=np.int64)
    if idx.ndim != 2 or idx.shape[1] != 4:
        raise ValueError("dihedral_indices must be shape (n, 4)")
    n = idx.shape[0]

    pbc_mode = pbc.lower()
    if pbc_mode not in ("none", "orthorhombic"):
        raise ValueError("pbc must be 'none' or 'orthorhombic'")
    use_plan = False
    if n == 1:
        try:
            from warp_md import DihedralPlan  # type: ignore

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

    if use_plan and DihedralPlan is not None and hasattr(system, "select_indices"):
        a_idx, b_idx, c_idx, d_idx = resolved[0]
        sel_a = system.select_indices([a_idx])
        sel_b = system.select_indices([b_idx])
        sel_c = system.select_indices([c_idx])
        sel_d = system.select_indices([d_idx])
        plan = DihedralPlan(sel_a, sel_b, sel_c, sel_d, False, pbc, False, False)
        angles = plan.run(traj, system, chunk_frames=chunk_frames, device=device)
        angles = np.asarray(angles, dtype=np.float64)
        if frame_indices is not None:
            n_frames = angles.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)
            angles = angles[select]
        angle = angles + phase
        cos = np.cos(angle)
        jvals = a * cos * cos + b * cos + c
        return jvals.reshape((-1, 1)).astype(np.float32)

    max_chunk = chunk_frames or 128
    chunk = traj.read_chunk(max_chunk)
    if chunk is None:
        raise ValueError("trajectory has no frames")
    values = []
    if frame_indices is not None:
        frame_set = {int(i) for i in frame_indices if int(i) >= 0}
    else:
        frame_set = None
    global_frame = 0

    while chunk is not None:
        coords = np.asarray(chunk["coords"], dtype=np.float64) * length_scale
        box = chunk.get("box")
        if box is not None:
            box = np.asarray(box, dtype=np.float64) * length_scale
        frames = coords.shape[0]
        for f in range(frames):
            if frame_set is not None and global_frame not in frame_set:
                global_frame += 1
                continue
            pos = coords[f]
            if pbc_mode == "orthorhombic":
                if box is None:
                    raise ValueError("pbc='orthorhombic' requires box in trajectory")
                b = box[f]
                if np.any(b == 0.0):
                    raise ValueError("box lengths must be non-zero for pbc")
            p0 = pos[idx[:, 0]]
            p1 = pos[idx[:, 1]]
            p2 = pos[idx[:, 2]]
            p3 = pos[idx[:, 3]]
            b0 = p0 - p1
            b1 = p2 - p1
            b2 = p3 - p2
            if pbc_mode == "orthorhombic":
                b0 = b0 - np.round(b0 / b) * b
                b1 = b1 - np.round(b1 / b) * b
                b2 = b2 - np.round(b2 / b) * b
            b1n = _axis_norm(b1)
            v = b0 - (b0 * b1n).sum(axis=1, keepdims=True) * b1n
            w = b2 - (b2 * b1n).sum(axis=1, keepdims=True) * b1n
            v = _axis_norm(v)
            w = _axis_norm(w)
            x = (v * w).sum(axis=1)
            y = (np.cross(b1n, v) * w).sum(axis=1)
            angle = np.arctan2(y, x) + phase
            cos = np.cos(angle)
            j = a * cos * cos + b * cos + c
            values.append(j.astype(np.float32))
            global_frame += 1
        chunk = traj.read_chunk(max_chunk)
    if values:
        values = np.vstack(values)
    else:
        values = np.empty((0, n), dtype=np.float32)
    return values


calc_ired_vector_and_matrix = ired_vector_and_matrix
calc_nh_order_parameters = nh_order_parameters

__all__ = [
    "jcoupling",
    "ired_vector_and_matrix",
    "nh_order_parameters",
    "calc_ired_vector_and_matrix",
    "calc_nh_order_parameters",
]
