lowess/

regression.rs

//! Local regression fitting for LOWESS smoothing.
//!
//! This module implements the local weighted least-squares regression used by
//! the HIGHER-level LOWESS pipeline. It provides a single-point fitter and
//! supporting utilities (weight computation, normalization, fallback policies,
//! and simple diagnostics) designed for robust production use.
//!
//! Global expectations
//! - Inputs x and y are numeric and aligned; many helpers assume x is sorted
//!   ascending. Debug-only assertions validate sorting but production code
//!   returns safe fallbacks for degenerate inputs.
//! - All numeric tolerances are conservative to avoid panics; callers may
//!   perform stricter validation upstream for performance or determinism.
//!
//! Primary parameters and flags
//! - x: `&[T]` — sorted (recommended) independent variable values.
//! - y: `&[T]` — dependent variable values aligned with x.
//! - idx: `usize` — index of the target point to fit (0..n-1).
//! - left/right: usize — inclusive window boundaries defining the local
//!   neighborhood used for the fit. These are clamped to [0, n-1] by helpers.
//! - use_robustness: bool — when true, per-observation robustness weights
//!   (from IRLS) are multiplied with kernel weights before normalization.
//! - robustness_weights: `&[T]` — per-observation multiplicative weights from a
//!   previous robustness pass. If not used, pass a slice of ones.
//! - weights: `&mut [T]` — scratch buffer for computed (unnormalized) weights;
//!   must be length n (or at least cover the positions accessed). The buffer
//!   is normalized in-place prior to regression.
//! - weight_fn: WeightFunction — kernel used to compute distance-based weights
//!   (Tricube, Epanechnikov, Gaussian, etc.). Bounded kernels support a fast
//!   short-circuit for |u| >= 1.
//! - zero_weight_fallback: ZeroWeightFallback — policy applied when the local
//!   sum of computed weights is zero. Options:
//!     * UseLocalMean — return the (unweighted) mean over `[left..=right]`.
//!     * ReturnOriginal — return `y[idx]`.
//!     * ReturnNone — propagate failure (caller decides).
//!
//! WeightParams specifics
//! - x_current: T — the x location being fitted.
//! - bandwidth: T — effective local half-width used for normalized distance u.
//!   Must be > 0 for a full regression; zero triggers constant-average fallback.
//! - h1: T — a tiny fraction of bandwidth below which kernel weight is forced to 1.
//!   This avoids numerical cancellation for extremely close points.
//! - h9: T — slightly less than bandwidth (e.g. 0.999*h) used to truncate the
//!   effective neighbor scan and determine the rightmost point to include.
//!
//! Functions and behaviors
//! - fit_point(ctx): primary entry. Computes kernel ± robustness weights,
//!   normalizes them, and runs a weighted linear least-squares fit evaluated at
//!   x_current. If weights sum to zero the configured fallback is used. If the
//!   weighted x-variance is too small, the fitter falls back to the weighted mean.
//! - compute_weights(...): fast, streaming weight computation that scans from
//!   left to right, short-circuits outside h9, applies h1 fast-path, and
//!   multiplies by robustness weights when requested. Returns the (unnormalized)
//!   total weight for the scanned region.
//! - find_rightmost_point(...): returns the largest index within h9 of x_current.
//! - normalize_weights(...): in-place normalization over [left..=right]. Debug
//!   builds assert sum > 0; production code expects callers to handle zero-sum.
//! - weighted_least_squares(...): numerically stable WLS for degree-1. Falls
//!   back to weighted average when denominator is below conservatively chosen
//!   tolerance (absolute and bandwidth-scaled relative terms).
//! - compute_weighted_average(...): assumes weights already normalized and
//!   returns ∑ wᵢ vᵢ over [left..=right].
//!
//! Debug & determinism
//! - Debug-only asserts check sorted x and buffer lengths; they do not change
//!   release behavior. Median/selection helpers used elsewhere prefer
//!   select_nth_unstable for performance (linear-time, not stable ordering).
//!
//! Production recommendations
//! - Pre-sort and deduplicate x/y upstream for reproducible window semantics.
//! - Provide a pre-allocated weights buffer to avoid repeated allocations.
//! - Choose an appropriate ZeroWeightFallback policy for your deployment:
//!   UseLocalMean for graceful smoothing, ReturnNone for strict failure modes.
//! - Use delta-driven interpolation and robust weights at the builder layer to
//!   control performance vs. robustness trade-offs for large datasets.

#[cfg(not(feature = "std"))]
extern crate alloc;
#[cfg(not(feature = "std"))]
use alloc::vec::Vec;

use crate::kernel::WeightFunction;
use num_traits::Float;

// ============================================================================
// Zero-weight fallback policy
// ============================================================================

/// Behavior to use when the computed total weight for a fit is zero.
#[derive(Copy, Clone, Debug, Default)]
pub enum ZeroWeightFallback {
    /// Fall back to the simple mean over the local window (current behavior).
    #[default]
    UseLocalMean,
    /// Return the original y value at the target index.
    ReturnOriginal,
    /// Return None (propagate failure) so caller can decide what to do.
    ReturnNone,
}

// ============================================================================
// Context Structures
// ============================================================================

/// Context for fitting a single point in LOWESS.
pub struct FitContext<'a, T> {
    /// X coordinates (sorted)
    pub x: &'a [T],
    /// Y coordinates
    pub y: &'a [T],
    /// Index of point to fit
    pub idx: usize,
    /// Left boundary of window
    pub left: usize,
    /// Right boundary of window
    pub right: usize,
    /// Whether to use robustness weights
    pub use_robustness: bool,
    /// Robustness weights from previous iteration
    pub robustness_weights: &'a [T],
    /// Scratch space for computed weights (length n)
    pub weights: &'a mut [T],
    /// Weight function to use
    pub weight_fn: WeightFunction,
    /// What to do if total weight is zero
    pub zero_weight_fallback: ZeroWeightFallback,
}

/// Parameters for weight computation.
pub struct WeightParams<T> {
    /// Current x value being fitted
    pub x_current: T,
    /// Bandwidth (window half-width)
    pub bandwidth: T,
    /// NOTE: h1 and h9 are implementation optimizations / truncation thresholds,
    /// not formal LOWESS parameters. They are used to fast-path near-zero
    /// distances (h1) and to truncate the kernel support for efficiency (h9).
    /// h1 = 0.001 * bandwidth ⇒ kernel weight forced to 1.0 for extremely
    /// close points; h9 = 0.999 * bandwidth used to stop scanning beyond the
    /// effective window in sorted x arrays.
    /// Low threshold (0.001 * bandwidth) – assign weight 1.0
    pub h1: T,
    /// High threshold (0.999 * bandwidth) – assign weight 0.0
    pub h9: T,
    /// Whether to use robustness weights
    pub use_robustness: bool,
    /// Weight function to use
    pub weight_fn: WeightFunction,
}

// ============================================================================
// Debug helpers
// ============================================================================

#[cfg(debug_assertions)]
fn assert_sorted<T: Float>(x: &[T]) {
    // Ensure x is non-decreasing. This is a debug-only check to document
    // and validate the sorted-x assumption made elsewhere (e.g. early break).
    for i in 1..x.len() {
        debug_assert!(
            x[i] >= x[i - 1],
            "x must be sorted non-decreasing (found x[{}] < x[{}])",
            i,
            i - 1
        );
    }
}

// ============================================================================
// Main Fitting Function
// ============================================================================

/// Fit a single point using local weighted regression.
///
/// # Returns
///
/// `Some(fitted_y)` on success, `None` if the weight sum is zero.
pub fn fit_point<T: Float>(ctx: FitContext<T>) -> Option<T> {
    let n = ctx.x.len();
    let x_current = ctx.x[ctx.idx];

    // ---- bandwidth -------------------------------------------------------
    let bandwidth = T::max(x_current - ctx.x[ctx.left], ctx.x[ctx.right] - x_current);
    if bandwidth <= T::zero() {
        // Degenerate window – fall back to weighted average of the whole window
        let mut sum_w = T::zero();
        for j in ctx.left..=ctx.right {
            sum_w = sum_w + ctx.weights[j];
        }
        if sum_w > T::zero() {
            return Some(compute_weighted_average(
                ctx.y,
                ctx.weights,
                ctx.left,
                ctx.right,
            ));
        } else {
            return Some(ctx.y[ctx.idx]);
        }
    }

    // ---- thresholds ------------------------------------------------------
    let h9 = T::from(0.999).unwrap() * bandwidth;
    let h1 = T::from(0.001).unwrap() * bandwidth;

    let params = WeightParams {
        x_current,
        bandwidth,
        h1,
        h9,
        use_robustness: ctx.use_robustness,
        weight_fn: ctx.weight_fn,
    };

    // ---- compute raw weights ---------------------------------------------
    let weight_sum = compute_weights(
        ctx.x,
        ctx.left,
        n,
        &params,
        ctx.use_robustness,
        ctx.robustness_weights,
        ctx.weights,
    );

    if weight_sum <= T::zero() {
        // Configurable fallback behavior when all computed weights are zero.
        match ctx.zero_weight_fallback {
            ZeroWeightFallback::UseLocalMean => {
                let cnt = T::from((ctx.right - ctx.left + 1) as f64).unwrap_or(T::one());
                let mean = ctx.y[ctx.left..=ctx.right]
                    .iter()
                    .copied()
                    .fold(T::zero(), |acc, v| acc + v)
                    / cnt;
                return Some(mean);
            }
            ZeroWeightFallback::ReturnOriginal => {
                return Some(ctx.y[ctx.idx]);
            }
            ZeroWeightFallback::ReturnNone => {
                return None;
            }
        }
    }

    // ---- effective rightmost point ---------------------------------------
    let right_most = find_rightmost_point(ctx.x, x_current, ctx.left, n, h9);

    // ---- normalize -------------------------------------------------------
    normalize_weights(ctx.weights, ctx.left, right_most, weight_sum);

    // ---- weighted least squares -------------------------------------------
    Some(weighted_least_squares(
        ctx.x,
        ctx.y,
        ctx.weights,
        ctx.left,
        right_most,
        x_current,
        bandwidth,
    ))
}

// ============================================================================
// Weight Computation
// ============================================================================

/// Compute kernel (and optional robustness) weights.
///
/// Uses the **fast** kernel path (`compute_weight_fast`) for speed.
pub fn compute_weights<T: Float>(
    x: &[T],
    left: usize,
    n: usize,
    params: &WeightParams<T>,
    use_robustness: bool,
    robustness_weights: &[T],
    weights: &mut [T],
) -> T {
    // Debug check: x must be sorted
    #[cfg(debug_assertions)]
    {
        assert_sorted(x);
    }

    // Degenerate bandwidth → zero weights
    if params.bandwidth <= T::zero() {
        for w in weights.iter_mut().take(n).skip(left) {
            *w = T::zero();
        }
        return T::zero();
    }

    let mut sum = T::zero();

    // First, skip all points to the left of (x_current - h9)
    let lower_bound = params.x_current - params.h9;
    let mut start = left;
    while start < n && x[start] < lower_bound {
        weights[start] = T::zero();
        start += 1;
    }

    // Then compute weights from start to right, breaking early if beyond h9
    for j in start..n {
        weights[j] = T::zero();

        let xj = x[j];
        let distance = (xj - params.x_current).abs();

        if distance > params.h9 {
            // Outside effective window
            if xj > params.x_current {
                break;
            }
            // To the left but outside h9, continue scanning
            continue;
        }

        // ---- kernel weight ------------------------------------------------
        let kernel = if distance <= params.h1 {
            T::one()
        } else {
            let u = distance / params.bandwidth;
            // *** fast path ***
            params.weight_fn.compute_weight_fast(u)
        };

        // ---- combine with robustness --------------------------------------
        let w = if use_robustness {
            kernel * robustness_weights[j]
        } else {
            kernel
        };

        weights[j] = w;
        sum = sum + w;
    }

    sum
}

/// Find the rightmost point whose distance ≤ `h9`.
pub fn find_rightmost_point<T: Float>(
    x: &[T],
    x_current: T,
    left: usize,
    n: usize,
    h9: T,
) -> usize {
    #[cfg(debug_assertions)]
    {
        assert_sorted(x);
    }

    let mut rightmost = left;
    for (j, &xj) in x.iter().enumerate().take(n).skip(left) {
        if (xj - x_current).abs() <= h9 {
            rightmost = j;
        } else {
            // If this point is to the right of x_current and outside h9 we can
            // break (future points will be farther). If it's to the left,
            // continue scanning since later points may fall within h9.
            if xj > x_current {
                break;
            }
            continue;
        }
    }
    rightmost
}

/// Normalize weights **in-place** over `[left, right]`.
///
/// # Panics
///
/// In debug builds if `sum == 0`.
#[inline]
pub fn normalize_weights<T: Float>(weights: &mut [T], left: usize, right: usize, sum: T) {
    debug_assert!(sum > T::zero(), "weight sum must be positive");
    let inv = T::one() / sum;
    for w in weights.iter_mut().take(right + 1).skip(left) {
        *w = *w * inv;
    }
}

// ============================================================================
// Weighted Least Squares
// ============================================================================

/// Weighted linear regression evaluated at `x_current`.
///
/// Falls back to a weighted average when the weighted variance is too small.
pub fn weighted_least_squares<T: Float>(
    x: &[T],
    y: &[T],
    weights: &[T],
    left: usize,
    right: usize,
    x_current: T,
    bandwidth: T,
) -> T {
    // Degenerate bandwidth → simple weighted average
    if bandwidth <= T::zero() {
        return compute_weighted_average(y, weights, left, right);
    }

    // ---- weighted means -------------------------------------------
    let x_mean = compute_weighted_average(x, weights, left, right);
    let y_mean = compute_weighted_average(y, weights, left, right);

    // ---- weighted variance (denominator) and covariance (numerator) ---
    let mut denom = T::zero();
    let mut numer = T::zero();
    for j in left..=right {
        let dx = x[j] - x_mean;
        denom = denom + weights[j] * dx * dx;
        numer = numer + weights[j] * dx * (y[j] - y_mean);
    }

    // ---- numerical stability check ------------------------------------
    // If the weighted x-variance (denom) is numerically tiny, fall back to
    // the weighted average. Use a hybrid tolerance: an absolute floor to
    // capture extremely narrow / near-collinear windows, and a tiny
    // bandwidth-scaled relative term to avoid underflow on very large scales.
    let abs_tol = T::from(1e-7).unwrap();
    let rel_tol = T::epsilon() * bandwidth * bandwidth;
    let tol = abs_tol.max(rel_tol);
    if denom <= tol {
        return compute_weighted_average(y, weights, left, right);
    }

    // ---- slope (β1) and fitted value ----------------------------------
    debug_assert!(denom > tol, "Division by near-zero denom");
    let beta1 = numer / denom;
    y_mean + beta1 * (x_current - x_mean)
}

/// Simple weighted average (∑ wᵢ yᵢ) – assumes weights already normalized.
#[inline]
pub fn compute_weighted_average<T: Float>(
    values: &[T],
    weights: &[T],
    left: usize,
    right: usize,
) -> T {
    let mut sum = T::zero();
    for j in left..=right {
        sum = sum + weights[j] * values[j];
    }
    sum
}

// ============================================================================
// Regression Diagnostics
// ============================================================================

/// Trace of the hat matrix – effective degrees of freedom.
pub fn compute_effective_df<T: Float>(weights_matrix: &[Vec<T>]) -> T {
    let n = weights_matrix.len();
    let mut trace = T::zero();

    for (i, row) in weights_matrix.iter().enumerate().take(n) {
        trace = trace + row[i];
    }

    trace
}

/// Leverage of a single observation (diagonal of the hat matrix).
#[inline]
pub fn compute_leverage<T: Float>(weights: &[T], idx: usize) -> T {
    weights[idx]
}

/// Residual variance estimate (weighted).
///
/// Uses `n_eff - 2` degrees of freedom (linear fit).
pub fn compute_residual_variance<T: Float>(
    residuals: &[T],
    weights: &[T],
    left: usize,
    right: usize,
) -> T {
    let mut sum_sq = T::zero();
    let mut n_eff = T::zero();

    for j in left..=right {
        if weights[j] > T::zero() {
            sum_sq = sum_sq + weights[j] * residuals[j] * residuals[j];
            n_eff = n_eff + T::one();
        }
    }

    let df = n_eff - T::from(2.0).unwrap();
    if df > T::zero() {
        sum_sq / df
    } else {
        T::zero()
    }
}

// ============================================================================
// Alternative Fitting Methods
// ============================================================================

/// Higher-degree polynomial fit (currently only degree 1 is implemented).
pub fn weighted_polynomial_fit<T: Float>(
    x: &[T],
    y: &[T],
    weights: &[T],
    left: usize,
    right: usize,
    x_current: T,
    degree: usize,
) -> T {
    if degree == 1 {
        let bandwidth = T::max(x_current - x[left], x[right] - x_current);
        weighted_least_squares(x, y, weights, left, right, x_current, bandwidth)
    } else {
        // For now just fall back to constant fit
        compute_weighted_average(y, weights, left, right)
    }
}

/// Locally constant (degree-0) regression – just a weighted average.
#[inline]
pub fn locally_constant_fit<T: Float>(y: &[T], weights: &[T], left: usize, right: usize) -> T {
    compute_weighted_average(y, weights, left, right)
}