"""Pure-numpy post-processing primitives for trajectory generation.

All functions are pure (no I/O, no global state) and accept an explicit
:class:`numpy.random.Generator` when randomness is involved.
"""
from __future__ import annotations

import numpy as np


def gaussian_smooth(x: np.ndarray, sigma: float = 1.0) -> np.ndarray:
    """5-tap gaussian smoothing along a 1-D array; endpoints preserved.

    Args:
        x: 1-D input array.
        sigma: gaussian std. Default 1.0 gives weights approximately
               [0.054, 0.244, 0.403, 0.244, 0.054].

    Returns:
        Smoothed array of the same shape. ``x[0]`` and ``x[-1]`` unchanged.
        If ``len(x) < 5`` returns a copy of ``x`` (kernel won't fit).
    """
    if len(x) < 5:
        return x.copy()
    kernel = np.exp(-0.5 * (np.arange(-2, 3) / sigma) ** 2)
    kernel /= kernel.sum()
    padded = np.pad(x, pad_width=2, mode="edge")
    smoothed = np.convolve(padded, kernel, mode="valid")
    smoothed[0] = x[0]
    smoothed[-1] = x[-1]
    return smoothed


def snap_endpoints(
    forward: np.ndarray,
    lateral: np.ndarray,
    seq_len: int,
    n_snap: int = 6,
) -> tuple[np.ndarray, np.ndarray]:
    """Force first point to (0,0) and last point to (1,0) with quadratic ease.

    The last ``n_snap`` points are linearly interpolated towards (1, 0)
    with quadratic easing, then the first/last points are pinned exactly.

    Args:
        forward: (T,) forward coordinates (modified in place).
        lateral: (T,) lateral coordinates (modified in place).
        seq_len: length of forward/lateral.
        n_snap: number of trailing points to ease (capped at seq_len//4).

    Returns:
        ``(forward, lateral)`` after modification.
    """
    n_snap = min(n_snap, seq_len // 4)
    for i in range(n_snap):
        alpha = ((i + 1) / n_snap) ** 2
        k = seq_len - n_snap + i
        forward[k] = forward[k] * (1.0 - alpha) + 1.0 * alpha
        lateral[k] = lateral[k] * (1.0 - alpha) + 0.0 * alpha
    forward[0], lateral[0] = 0.0, 0.0
    forward[-1], lateral[-1] = 1.0, 0.0
    return forward, lateral


def smooth_start(
    forward: np.ndarray,
    lateral: np.ndarray,
    n: int = 4,
) -> tuple[np.ndarray, np.ndarray]:
    """Dampen lateral oscillation in the first ``n`` points.

    Assumes :func:`snap_endpoints` has already pinned (0,0). Forward is
    forced non-decreasing locally; lateral is linearly damped towards 0.
    """
    n_start_fix = min(n, len(forward) // 4)
    for i in range(1, n_start_fix + 1):
        blend = i / (n_start_fix + 1)
        forward[i] = max(forward[i], forward[i - 1])
        lateral[i] = lateral[i] * blend
    return forward, lateral


def enforce_forward_monotonic(forward: np.ndarray) -> np.ndarray:
    """Force ``forward`` non-decreasing, clip to [0,1], pin endpoints."""
    seq_len = len(forward)
    for i in range(1, seq_len - 1):
        if forward[i] < forward[i - 1]:
            forward[i] = forward[i - 1] + 0.001
    forward = np.clip(forward, 0.0, 1.0)
    forward[0] = 0.0
    forward[-1] = 1.0
    return forward


def resample_arc(xy: np.ndarray, n_points: int) -> np.ndarray:
    """Resample a 2-D polyline to ``n_points`` along cumulative arc length."""
    arc = np.concatenate(
        [[0], np.cumsum(np.linalg.norm(np.diff(xy, axis=0), axis=1))]
    )
    s_new = np.linspace(0, arc[-1], n_points)
    return np.stack(
        [np.interp(s_new, arc, xy[:, 0]), np.interp(s_new, arc, xy[:, 1])],
        axis=1,
    )


def build_timestamps(
    log_dt: np.ndarray,
    total_duration_ms: float,
    dt_clip: tuple[float, float] = (2.0, 150.0),
) -> np.ndarray:
    """Convert per-step log_dt + total duration to cumulative ms timestamps.

    Args:
        log_dt: (N,) array of natural-log step intervals.
        total_duration_ms: target total span. The output is scaled so the
            sum approximately matches this (modulo dt_clip).
        dt_clip: (min, max) per-step clamp in milliseconds.

    Returns:
        (N,) integer-rounded cumulative timestamps starting at 0,
        strictly increasing.
    """
    n = len(log_dt)
    dt_raw = np.clip(np.exp(log_dt), 0.0, None)
    dt_sum = dt_raw.sum()
    if dt_sum > 1e-6:
        scale = total_duration_ms / dt_sum
    else:
        scale = total_duration_ms / max(n, 1)
    dt_ms = np.clip(dt_raw * scale, dt_clip[0], dt_clip[1])

    t_abs = np.cumsum(dt_ms)
    t_abs = np.concatenate([[0.0], t_abs[:-1]])

    for i in range(1, n):
        if t_abs[i] <= t_abs[i - 1]:
            t_abs[i] = t_abs[i - 1] + 1.0
    return t_abs


def truncnorm_sample(
    mu: float,
    sigma: float,
    low: float,
    high: float,
    rng: np.random.Generator,
    max_tries: int = 32,
) -> float:
    """Sample from N(mu, sigma) truncated to [low, high] via rejection.

    Falls back to clipping if rejection fails ``max_tries`` times.
    """
    for _ in range(max_tries):
        v = rng.normal(mu, sigma)
        if low <= v <= high:
            return float(v)
    return float(np.clip(rng.normal(mu, sigma), low, high))


def sample_duration(
    duration_dist: dict,
    dist: float,
    rng: np.random.Generator,
) -> float:
    """Sample total trajectory duration (ms) for the given pixel distance.

    Uses per-bin log-normal parameters in ``duration_dist``.
    """
    bins = duration_dist["bins"]
    params = duration_dist["params"]
    bin_idx = len(bins) - 1
    for i in range(len(bins) - 1):
        if dist < bins[i + 1]:
            bin_idx = i
            break
    bin_idx = min(bin_idx, len(params) - 1)
    mu_log = params[bin_idx]["mu_log"]
    sigma_log = params[bin_idx]["sigma_log"]
    return float(np.exp(rng.normal(mu_log, sigma_log)))
