# Usage:
# from warp_md.analysis.diffusion import diffusion
# out = diffusion(traj, system, mask=':WAT@O', tstep=1.0)

from __future__ import annotations

from typing import Optional, Sequence, Union

import numpy as np

from .geometry import _dihedral_from_points


def _all_resid_mask(system) -> str:
    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: str) -> np.ndarray:
    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 = []
    time_list = []
    chunk = traj.read_chunk(max_chunk)
    if chunk is None:
        return None, None
    while chunk is not None:
        coords_list.append(np.asarray(chunk["coords"], dtype=np.float64))
        time = chunk.get("time_ps")
        if time is None:
            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))
    times = np.concatenate(time_list, axis=0) if time_list else None
    return coords, times


def _apply_frame_indices(coords, times, frame_indices: Optional[Sequence[int]]):
    if frame_indices is None:
        return coords, times
    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 times is not None:
        times = times[select]
    return coords, times


def _lag_time(times: Optional[np.ndarray], tstep: float, n_lags: int) -> np.ndarray:
    if times is None or times.size < 2:
        return np.arange(n_lags, dtype=np.float64) * float(tstep)
    dt = np.diff(times)
    dt = dt[np.isfinite(dt) & (dt > 0.0)]
    if dt.size == 0:
        return np.arange(n_lags, dtype=np.float64) * float(tstep)
    step = float(np.median(dt))
    return np.arange(n_lags, dtype=np.float64) * step


def diffusion(
    traj,
    system,
    mask: str = "",
    tstep: float = 1.0,
    individual: bool = False,
    frame_indices: Optional[Sequence[int]] = None,
    chunk_frames: Optional[int] = None,
):
    """Compute diffusion from mean-squared displacement.

    Returns dict with time, MSD components, and diffusion estimate.
    """
    if tstep <= 0.0:
        raise ValueError("tstep must be positive")

    coords, times = _read_all(traj, chunk_frames)
    if coords is None:
        raise ValueError("trajectory has no frames")
    if coords.size == 0:
        empty = np.empty((0,), dtype=np.float32)
        return {"time": empty, "X": empty, "Y": empty, "Z": empty, "MSD": empty, "D": 0.0}

    coords, times = _apply_frame_indices(coords, times, frame_indices)
    if coords.size == 0:
        empty = np.empty((0,), dtype=np.float32)
        return {"time": empty, "X": empty, "Y": empty, "Z": empty, "MSD": empty, "D": 0.0}

    idx = _selection_indices(system, mask)
    if idx.size == 0:
        empty = np.empty((0,), dtype=np.float32)
        return {"time": empty, "X": empty, "Y": empty, "Z": empty, "MSD": empty, "D": 0.0}

    sel = coords[:, idx, :]
    n_frames = sel.shape[0]
    n_lags = n_frames

    msd_x = np.zeros(n_lags, dtype=np.float64)
    msd_y = np.zeros(n_lags, dtype=np.float64)
    msd_z = np.zeros(n_lags, dtype=np.float64)
    msd = np.zeros(n_lags, dtype=np.float64)
    counts = np.zeros(n_lags, dtype=np.int64)

    if individual:
        msd_individual = np.zeros((idx.size, n_lags), dtype=np.float64)
    else:
        msd_individual = None

    for lag in range(n_lags):
        disp = sel[lag:] - sel[: n_frames - lag]
        if disp.size == 0:
            continue
        sq = disp * disp
        msd_x[lag] = np.mean(sq[..., 0])
        msd_y[lag] = np.mean(sq[..., 1])
        msd_z[lag] = np.mean(sq[..., 2])
        msd[lag] = np.mean(sq.sum(axis=2))
        counts[lag] = disp.shape[0]
        if individual:
            msd_individual[:, lag] = np.mean(sq.sum(axis=2), axis=0)

    time = _lag_time(times, tstep, n_lags)

    def _linear_fit(x, y):
        if x.size < 2:
            return 0.0
        x = np.asarray(x, dtype=np.float64)
        y = np.asarray(y, dtype=np.float64)
        A = np.vstack([x, np.ones_like(x)]).T
        slope, _ = np.linalg.lstsq(A, y, rcond=None)[0]
        return float(slope)

    start = max(1, n_lags // 2)
    slope = _linear_fit(time[start:], msd[start:])
    d_total = slope / 6.0 if slope > 0.0 else 0.0

    out = {
        "time": time.astype(np.float32),
        "X": msd_x.astype(np.float32),
        "Y": msd_y.astype(np.float32),
        "Z": msd_z.astype(np.float32),
        "MSD": msd.astype(np.float32),
        "D": float(d_total),
    }
    if individual and msd_individual is not None:
        out["MSD_individual"] = msd_individual.astype(np.float32)
    return out


def _torsion_stats(coords: np.ndarray, torsions: np.ndarray) -> np.ndarray:
    if torsions.size == 0:
        return np.zeros((coords.shape[0], 4), dtype=np.float32)
    if torsions.size % 4 != 0:
        raise ValueError("torsion indices length must be a multiple of 4")
    n_torsions = torsions.size // 4
    out = np.zeros((coords.shape[0], 4), dtype=np.float32)
    for f in range(coords.shape[0]):
        trans = cis = g_plus = g_minus = 0
        for t in range(n_torsions):
            a, b, c, d = torsions[t * 4 : t * 4 + 4]
            ang = _dihedral_from_points(
                coords[f : f + 1, a, :],
                coords[f : f + 1, b, :],
                coords[f : f + 1, c, :],
                coords[f : f + 1, d, :],
            )[0]
            if ang > 150.0 or ang < -150.0:
                trans += 1
            elif abs(ang) < 30.0:
                cis += 1
            elif ang > 0.0:
                g_plus += 1
            else:
                g_minus += 1
        denom = float(max(n_torsions, 1))
        out[f, 0] = trans / denom
        out[f, 1] = cis / denom
        out[f, 2] = g_plus / denom
        out[f, 3] = g_minus / denom
    return out


def _torsion_stats_weighted(coords: np.ndarray, torsions: np.ndarray, masses: np.ndarray) -> np.ndarray:
    if torsions.size == 0:
        return np.zeros((coords.shape[0], 4), dtype=np.float32)
    if torsions.size % 4 != 0:
        raise ValueError("torsion indices length must be a multiple of 4")
    n_torsions = torsions.size // 4
    out = np.zeros((coords.shape[0], 4), dtype=np.float32)
    for f in range(coords.shape[0]):
        trans = cis = g_plus = g_minus = 0.0
        total = 0.0
        for t in range(n_torsions):
            a, b, c, d = torsions[t * 4 : t * 4 + 4]
            ang = _dihedral_from_points(
                coords[f : f + 1, a, :],
                coords[f : f + 1, b, :],
                coords[f : f + 1, c, :],
                coords[f : f + 1, d, :],
            )[0]
            weight = 1.0
            if masses.size > 0:
                weight = float(
                    (masses[a] + masses[b] + masses[c] + masses[d]) / 4.0
                )
            if ang > 150.0 or ang < -150.0:
                trans += weight
            elif abs(ang) < 30.0:
                cis += weight
            elif ang > 0.0:
                g_plus += weight
            else:
                g_minus += weight
            total += weight
        denom = total if total > 0.0 else 1.0
        out[f, 0] = trans / denom
        out[f, 1] = cis / denom
        out[f, 2] = g_plus / denom
        out[f, 3] = g_minus / denom
    return out


def tordiff(
    traj,
    system,
    mask: Union[str, Sequence[int], np.ndarray] = "",
    mass: bool = False,
    out: Optional[str] = None,
    diffout: Optional[str] = None,
    time: float = 1.0,
    extra_options: str = "",
    frame_indices: Optional[Sequence[int]] = None,
    chunk_frames: Optional[int] = None,
    device: str = "auto",
):
    """Toroidal diffusion (torsion state fractions)."""
    del extra_options
    coords, _times = _read_all(traj, chunk_frames)
    if coords is None:
        raise ValueError("trajectory has no frames")
    if coords.size == 0:
        empty = np.zeros((0, 4), dtype=np.float32)
        return {
            "time": np.zeros((0,), dtype=np.float32),
            "trans": empty[:, 0],
            "cis": empty[:, 1],
            "g_plus": empty[:, 2],
            "g_minus": empty[:, 3],
        }

    coords, _times = _apply_frame_indices(coords, _times, frame_indices)
    if coords.size == 0:
        empty = np.zeros((0, 4), dtype=np.float32)
        return {
            "time": np.zeros((0,), dtype=np.float32),
            "trans": empty[:, 0],
            "cis": empty[:, 1],
            "g_plus": empty[:, 2],
            "g_minus": empty[:, 3],
        }

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

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

    if mass:
        use_plan = False

    if use_plan and TorsionDiffusionPlan is not None:
        if isinstance(mask, str):
            sel = system.select(mask)
        else:
            indices = np.asarray(mask, dtype=np.int64).reshape(-1).tolist()
            sel = system.select_indices(indices)
        plan = TorsionDiffusionPlan(sel)
        from .trajectory import ArrayTrajectory

        temp = ArrayTrajectory(coords.astype(np.float32))
        mat = plan.run(temp, system, chunk_frames=None, device=device)
        data = np.asarray(mat, dtype=np.float32)
    else:
        if isinstance(mask, str):
            sel = system.select(mask)
            torsions = np.asarray(list(sel.indices), dtype=np.int64)
        else:
            torsions = np.asarray(mask, dtype=np.int64).reshape(-1)
        if mass:
            if hasattr(system, "atom_table"):
                atoms_tbl = system.atom_table()
                masses = np.asarray(atoms_tbl.get("mass", []), dtype=np.float64)
            else:
                masses = np.asarray([], dtype=np.float64)
            data = _torsion_stats_weighted(coords, torsions, masses)
        else:
            data = _torsion_stats(coords, torsions)

    n_frames = data.shape[0]
    if _times is not None and _times.size == n_frames:
        time_arr = _times.astype(np.float32)
    else:
        time_arr = np.arange(n_frames, dtype=np.float32) * float(time)

    out_dict = {
        "time": time_arr,
        "trans": data[:, 0],
        "cis": data[:, 1],
        "g_plus": data[:, 2],
        "g_minus": data[:, 3],
    }

    if out:
        table = np.column_stack([time_arr, data[:, 0], data[:, 1], data[:, 2], data[:, 3]])
        np.savetxt(out, table, header="time trans cis g_plus g_minus")
    if diffout:
        if n_frames < 2:
            diff_table = np.empty((0, 5), dtype=np.float32)
        else:
            dt = np.diff(time_arr)
            dt = np.where(dt == 0.0, 1.0, dt)
            diff = (data[1:] - data[:-1]) / dt[:, None]
            diff_table = np.column_stack([time_arr[1:], diff[:, 0], diff[:, 1], diff[:, 2], diff[:, 3]])
        np.savetxt(diffout, diff_table, header="time dtrans dcis dg_plus dg_minus")

    return out_dict


def toroidal_diffusion(*args, **kwargs):
    return tordiff(*args, **kwargs)


__all__ = ["diffusion", "tordiff", "toroidal_diffusion"]
