"""Kinematic metrics for mouse trajectory evaluation.

All inputs are 1-D NumPy arrays. Time is in milliseconds, position in pixels.
Velocities are px/ms, accelerations px/ms², jerks px/ms³.
"""
from __future__ import annotations

import numpy as np


def compute_speed(
    xs: np.ndarray, ys: np.ndarray, ts: np.ndarray, eps: float = 1e-6
) -> np.ndarray:
    """Compute scalar speed at each step.

    Args:
        xs: (N,) x coordinates.
        ys: (N,) y coordinates.
        ts: (N,) timestamps in ms.
        eps: minimum dt (ms) to avoid div-by-zero.

    Returns:
        (N-1,) array of speeds (px/ms).
    """
    dx = np.diff(xs)
    dy = np.diff(ys)
    dt = np.maximum(np.diff(ts), eps)
    return np.hypot(dx, dy) / dt


def compute_acceleration(speeds: np.ndarray, ts: np.ndarray, eps: float = 1e-6) -> np.ndarray:
    """Compute scalar acceleration from speeds.

    Args:
        speeds: (M,) speeds (px/ms). Typically M = N-1 from compute_speed.
        ts: (M,) or (M+1,) timestamps in ms.
            If len(ts) == len(speeds): timestamps are treated as the time
            points associated with each speed value directly.
            If len(ts) == len(speeds)+1: timestamps are the original position
            timestamps; midpoints are computed for speed intervals.
        eps: minimum dt (ms) to avoid div-by-zero.

    Returns:
        (M-1,) array of accelerations (px/ms²).
    """
    if len(speeds) < 2:
        return np.array([], dtype=float)
    if len(ts) == len(speeds):
        # ts[i] is already the time associated with speed[i]
        dt = np.maximum(np.diff(ts), eps)
    else:
        # ts has length M+1; speed[i] is between ts[i] and ts[i+1]
        midpoints = (ts[:-1] + ts[1:]) / 2.0
        dt = np.maximum(np.diff(midpoints), eps)
    return np.diff(speeds) / dt


def compute_jerk(accels: np.ndarray, ts: np.ndarray, eps: float = 1e-6) -> np.ndarray:
    """Compute jerk from accelerations.

    Args:
        accels: (K,) accelerations.
        ts: (K+2,) timestamps that produced those accelerations.
            Used to derive midpoint-of-midpoint dts.
        eps: minimum dt to avoid div-by-zero.

    Returns:
        (K-1,) array of jerks (px/ms³).
    """
    if len(accels) < 2:
        return np.array([], dtype=float)
    # Approximate dt for jerks as average dt of original ts (good enough for stats)
    dt_avg = np.maximum(np.diff(ts).mean(), eps)
    return np.diff(accels) / dt_avg


def compute_stats(x: np.ndarray) -> dict[str, float]:
    """Summary statistics for a 1-D distribution.

    Returns:
        dict with keys: mean, std, cv (coef of variation), p25, p50, p75, p95.
    """
    if len(x) == 0:
        return {k: 0.0 for k in ("mean", "std", "cv", "p25", "p50", "p75", "p95")}
    x = np.asarray(x, dtype=float)
    mean = float(x.mean())
    std = float(x.std(ddof=1)) if len(x) > 1 else 0.0
    cv = std / mean if mean != 0 else 0.0
    return {
        "mean": mean,
        "std": std,
        "cv": cv,
        "p25": float(np.percentile(x, 25)),
        "p50": float(np.percentile(x, 50)),
        "p75": float(np.percentile(x, 75)),
        "p95": float(np.percentile(x, 95)),
    }


def fft_spectrum(
    signal: np.ndarray, sample_rate_hz: float
) -> tuple[np.ndarray, np.ndarray]:
    """Compute one-sided FFT magnitude spectrum.

    Args:
        signal: 1-D real-valued signal.
        sample_rate_hz: Sampling rate in Hz.

    Returns:
        (freqs, magnitudes) — positive frequencies only.
        Magnitudes are absolute values of complex FFT coefficients.
    """
    n = len(signal)
    if n == 0:
        return np.array([]), np.array([])
    # Zero-mean to remove DC component which dominates the spectrum
    s = signal - signal.mean()
    fft = np.fft.rfft(s)
    freqs = np.fft.rfftfreq(n, d=1.0 / sample_rate_hz)
    return freqs, np.abs(fft)


def kl_divergence_histograms(
    x: np.ndarray,
    y: np.ndarray,
    bins: int = 50,
    eps: float = 1e-10,
) -> float:
    """KL divergence KL(P_x || P_y) estimated via shared-bin histograms.

    Both arrays are histogrammed over their joint range. Empty bins get
    `eps` mass to avoid log(0) — keeps result finite even for disjoint
    supports.

    Args:
        x: samples from distribution P.
        y: samples from distribution Q (the "reference").
        bins: number of histogram bins.
        eps: smoothing constant for empty bins.

    Returns:
        scalar KL divergence (nats). Always finite, ≥ 0.
    """
    if len(x) == 0 or len(y) == 0:
        return 0.0
    lo = float(min(x.min(), y.min()))
    hi = float(max(x.max(), y.max()))
    if hi <= lo:
        return 0.0
    edges = np.linspace(lo, hi, bins + 1)
    px, _ = np.histogram(x, bins=edges, density=False)
    qy, _ = np.histogram(y, bins=edges, density=False)
    px = px.astype(float) + eps
    qy = qy.astype(float) + eps
    px /= px.sum()
    qy /= qy.sum()
    return float(np.sum(px * np.log(px / qy)))