"""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 damp_start(lateral: np.ndarray, n: int = 4) -> np.ndarray:
    """Dampen lateral oscillation over the first ``n`` points, continuously.

    Weights follow smoothstep(i / (n+1)) so the damping releases smoothly
    into the untouched region (the old linear blend jumped from 0.8 to
    1.0 and left a visible kink).

    Args:
        lateral: (T,) lateral coordinates.
        n: number of leading points to dampen (capped at len//4).

    Returns:
        New (T,) array.
    """
    out = lateral.copy()
    n = min(n, len(out) // 4)
    for i in range(1, n + 1):
        t = i / (n + 1)
        w = t * t * (3.0 - 2.0 * t)  # smoothstep
        out[i] *= w
    return out


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)))


def soften_forward(
    forward: np.ndarray,
    backtrack_tol: float = 0.02,
    overshoot_span: float = 0.08,
) -> np.ndarray:
    """Softly regularise the forward axis without destroying natural motion.

    Real trajectories contain small backward corrections and overshoot
    past the target; hard clipping turns both into visible artifacts
    (stacked points, vertical walls). Instead:

    - Backtracking is tolerated up to ``backtrack_tol``; larger dips are
      raised to ``prev - backtrack_tol``.
    - Values above 1.0 are compressed with tanh so they asymptote at
      ``1 + overshoot_span`` (moderate overshoot survives, extremes
      cannot fly far past the target).
    - No lower clip: the endpoint warp pins the first point later.

    Args:
        forward: (T,) forward coordinates.
        backtrack_tol: max allowed per-step regression.
        overshoot_span: asymptotic max excess above 1.0.

    Returns:
        New (T,) array.
    """
    out = forward.copy()
    for i in range(1, len(out)):
        floor = out[i - 1] - backtrack_tol
        if out[i] < floor:
            out[i] = floor
    over = out > 1.0
    out[over] = 1.0 + overshoot_span * np.tanh((out[over] - 1.0) / overshoot_span)
    return out


def warp_endpoints(
    forward: np.ndarray,
    lateral: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
    """Warp the whole curve so endpoints land exactly on (0,0) and (1,0).

    Computes the residual of the first point vs (0, 0) and the last point
    vs (1, 0), then subtracts each residual weighted by a smoothstep that
    is 1 at its own end and 0 at the opposite end. Unlike tail-dragging,
    this preserves the trajectory's local shape and approach direction.

    Args:
        forward: (T,) forward coordinates.
        lateral: (T,) lateral coordinates.

    Returns:
        ``(forward, lateral)`` new arrays, endpoints pinned exactly.
    """
    t = np.linspace(0.0, 1.0, len(forward))
    w_end = t * t * (3.0 - 2.0 * t)      # smoothstep: 0 at start → 1 at end
    w_start = 1.0 - w_end                # mirrored

    res_f0, res_l0 = forward[0] - 0.0, lateral[0] - 0.0
    res_f1, res_l1 = forward[-1] - 1.0, lateral[-1] - 0.0

    fo = forward - w_start * res_f0 - w_end * res_f1
    lo = lateral - w_start * res_l0 - w_end * res_l1

    fo[0], lo[0] = 0.0, 0.0
    fo[-1], lo[-1] = 1.0, 0.0
    return fo, lo
