//! Local weighted regression for LOWESS smoothing.
//!
//! ## Purpose
//!
//! This module provides the core fitting algorithm for LOWESS: local weighted
//! linear regression. At each point, a linear model is fit using nearby points
//! with weights determined by a kernel function. This is the heart of the
//! LOWESS smoothing process.
//!
//! ## Design notes
//!
//! * Uses weighted least squares (WLS) for local linear regression.
//! * Weights are computed from kernel functions (tricube, Epanechnikov, etc.).
//! * Optionally combines kernel weights with robustness weights for outlier resistance.
//! * Implements zero-weight fallback policies for numerical stability.
//! * All fitting operations are pure functions with explicit parameters (no hidden state).
//! * Generic over `Float` types to support f32 and f64.
//!
//! ## Key concepts
//!
//! ### Local Weighted Regression
//! For each point x_i:
//! 1. Define neighborhood window (nearest k points)
//! 2. Compute window radius (max distance to window boundary)
//! 3. Compute kernel weights: w_j = K(|x_j - x_i| / h)
//! 4. Optionally combine with robustness weights: w_j ← w_j × r_j
//! 5. Fit weighted linear model: ŷ = β₀ + β₁(x - x̄)
//! 6. Evaluate at x_i to get smoothed value
//!
//! ### Weighted Least Squares
//! Minimizes the weighted sum of squared residuals:
//! ```text
//! min Σ w_j (y_j - β₀ - β₁x_j)²
//! ```
//!
//! Solution:
//! ```text
//! β₁ = Σ w_j (x_j - x̄)(y_j - ȳ) / Σ w_j (x_j - x̄)²
//! β₀ = ȳ - β₁ x̄
//! ```
//!
//! where x̄ and ȳ are weighted means.
//!
//! ### Robustness Weights
//! In iterative robustness refinement, points with large residuals receive
//! lower weights. These robustness weights are multiplied with kernel weights
//! to downweight outliers.
//!
//! ### Zero-Weight Fallback
//! When all weights are zero (e.g., very small window radius or all robustness
//! weights zero), fallback policies determine the behavior:
//! - Use local mean (default)
//! - Return original value
//! - Return None (propagate failure)
//!
//! ## Invariants
//!
//! * Window radius must be positive for meaningful regression.
//! * Weights are normalized to sum to 1 before fitting.
//! * Fitted values are always finite (no NaN or infinity).
//! * Window indices are always within array bounds.
//!
//! ## Non-goals
//!
//! * This module does not compute the windows (done by window module).
//! * This module does not compute robustness weights (done by robustness module).
//! * This module does not perform higher-degree polynomial regression.
//! * This module does not validate input data for sorting or duplicates.
//!
//! ## Visibility
//!
//! The [`Regression`] trait and [`LinearRegression`] struct are part of the
//! internal API used by the LOWESS engine. The [`RegressionContext`] struct
//! encapsulates all parameters for fitting. These may change without notice.

use crate::primitives::window::Window;
use core::fmt;
use num_traits::Float;

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

use crate::math::kernel::WeightFunction;

// ============================================================================
// Regression Context
// ============================================================================

/// Context containing all data needed to fit a single point.
///
/// This struct allows regression implementations to be pure functions that
/// take all their inputs explicitly, with no hidden state. This design
/// facilitates testing, parallelization, and reasoning about the code.
///
/// # Fields
///
/// * `x` - Independent variable values (sorted)
/// * `y` - Dependent variable values
/// * `idx` - Index of the point to fit
/// * `window` - Neighborhood defining local regression region
/// * `use_robustness` - Whether to incorporate robustness weights
/// * `robustness_weights` - Per-point robustness weights (all 1.0 if not using)
/// * `weights` - Working buffer for computed weights (will be filled)
/// * `weight_function` - Kernel function for distance-based weighting
/// * `zero_weight_fallback` - Policy for handling all-zero weights
pub struct RegressionContext<'a, T: Float> {
    /// Slice of x-values (independent variable)
    pub x: &'a [T],

    /// Slice of y-values (dependent variable)
    pub y: &'a [T],

    /// Index of the point to fit
    pub idx: usize,

    /// Window for the local fit (defines neighborhood)
    pub window: Window,

    /// Whether to use robustness weights
    ///
    /// When true, combine kernel weights with robustness weights to
    /// downweight outliers.
    pub use_robustness: bool,

    /// Slice of robustness weights (all 1.0 if not using robustness)
    pub robustness_weights: &'a [T],

    /// Mutable slice of weights to be used in fitting
    ///
    /// This is a working buffer that will be filled with the product of
    /// kernel weights and (optionally) robustness weights.
    pub weights: &'a mut [T],

    /// Weight function (kernel)
    ///
    /// Enum specifying the kernel to use for distance-based weighting.
    pub weight_function: WeightFunction,

    /// Zero-weight fallback policy
    ///
    /// Defines what to do if all computed weights are zero.
    pub zero_weight_fallback: ZeroWeightFallback,
}

// ============================================================================
// Regression Trait
// ============================================================================

/// Trait for fitting a local regression at a single point.
///
/// Implementations of this trait define how to fit a point given a context
/// with data, weights, and parameters. The trait is `Sync + Send` to support
/// parallel implementations.
///
/// # Typical Implementations
///
/// * **LinearRegression**: Locally weighted linear regression (degree 1)
/// * **ConstantFitter**: Locally weighted constant (mean) fit (degree 0)
/// * **QuadraticFitter**: Locally weighted quadratic regression (degree 2)
pub trait Regression<T: Float>: Sync + Send {
    /// Fit the point specified in the context.
    ///
    /// # Parameters
    ///
    /// * `context` - Context containing all data and parameters for the fit
    ///
    /// # Returns
    ///
    /// * `Some(fitted_value)` - If the fit was successful
    /// * `None` - If the fit could not be computed (e.g., singular matrix,
    ///   not enough points, or zero-weight fallback policy returns None)
    fn fit(&self, context: RegressionContext<T>) -> Option<T>;
}

// ============================================================================
// Linear Regression Implementation
// ============================================================================

/// Standard local linear regression fitter.
///
/// This struct implements the [`Regression`] trait using standard local
/// weighted linear regression (degree 1). This is the default fitting
/// strategy used by LOWESS and follows Cleveland's original implementation.
///
/// # Algorithm
///
/// At each point `x_i` with window `[left, right]`:
/// 1. Compute window radius: `h = max(x_i - x[left], x[right] - x_i)`
/// 2. Compute kernel weights: `w_j = K(|x_j - x_i| / h)` for `j ∈ [left, right]`
/// 3. Optionally apply robustness weights: `w_j ← w_j × r_j`
/// 4. Normalize weights: `w_j ← w_j / Σw_k`
/// 5. Fit weighted line: `ŷ = ȳ + β₁(x - x̄)`
///    where `β₁ = Σ w_j (x_j - x̄)(y_j - ȳ) / Σ w_j (x_j - x̄)²`
/// 6. Evaluate at `x_i`
///
/// # Special cases
///
/// * **Zero window radius**: Falls back to weighted average
/// * **Zero weight sum**: Applies zero-weight fallback policy
/// * **Zero variance**: Returns horizontal line at weighted mean
#[derive(Debug, Clone, Copy, Default)]
pub struct LinearRegression;

impl<T: Float> Regression<T> for LinearRegression {
    fn fit(&self, context: RegressionContext<T>) -> Option<T> {
        let n = context.x.len();

        // Defensive bounds check
        if context.idx >= n {
            return None;
        }

        // Validate window bounds
        if context.window.left >= n || context.window.right >= n {
            return None;
        }

        let x_current = context.x[context.idx];

        // Compute window radius
        let window_radius = context.window.max_distance(context.x, x_current);

        if window_radius <= T::zero() {
            // Degenerate window (all x-values identical): fall back to weighted average
            let mut sum_w = T::zero();
            let mut sum_wy = T::zero();
            for j in context.window.left..=context.window.right {
                let w = context.weights[j];
                sum_w = sum_w + w;
                sum_wy = sum_wy + w * context.y[j];
            }
            if sum_w > T::zero() {
                return Some(sum_wy);
            } else {
                return Some(context.y[context.idx]);
            }
        }

        let weight_params = WeightParams::new(x_current, window_radius, context.use_robustness);

        // Compute raw kernel weights (unnormalized, without robustness)
        let (mut weight_sum, rightmost_idx) = context.weight_function.compute_window_weights(
            context.x,
            context.window.left,
            context.window.right,
            weight_params.x_current,
            weight_params.window_radius,
            weight_params.h1,
            weight_params.h9,
            context.weights,
        );

        // Apply robustness weights if needed
        if context.use_robustness {
            // Re-calculate sum after applying robustness weights
            weight_sum = T::zero();
            for j in context.window.left..=rightmost_idx {
                if context.weights[j] > T::zero() {
                    let w_robust = context.robustness_weights[j];
                    let w_final = context.weights[j] * w_robust;
                    context.weights[j] = w_final;
                    weight_sum = weight_sum + w_final;
                }
            }
        }

        if weight_sum <= T::zero() {
            // Handle zero-weight case according to fallback policy
            match context.zero_weight_fallback {
                ZeroWeightFallback::UseLocalMean => {
                    let window_size = context
                        .window
                        .right
                        .saturating_sub(context.window.left)
                        .saturating_add(1);
                    let cnt = T::from(window_size as f64).unwrap_or(T::one());
                    let mean = context.y[context.window.left..=context.window.right]
                        .iter()
                        .copied()
                        .fold(T::zero(), |acc, v| acc + v)
                        / cnt;
                    return Some(mean);
                }
                ZeroWeightFallback::ReturnOriginal => {
                    return Some(context.y[context.idx]);
                }
                ZeroWeightFallback::ReturnNone => {
                    return None;
                }
            }
        }

        // Normalize weights to sum to 1
        debug_assert!(
            weight_sum > T::zero(),
            "normalize_range: sum must be positive"
        );
        let inv_sum = T::one() / weight_sum;
        for j in context.window.left..=rightmost_idx {
            context.weights[j] = context.weights[j] * inv_sum;
        }

        // Perform weighted least squares regression
        Some(GLSModel::local_wls(
            context.x,
            context.y,
            context.weights,
            context.window.left,
            rightmost_idx,
            x_current,
            window_radius,
        ))
    }
}

// ============================================================================
// Zero-Weight Fallback Policy
// ============================================================================

/// Policy for handling cases where all weights are zero.
///
/// This enum defines the behavior when the sum of weights for a point is zero,
/// which can happen if the window radius is very small or if robustness weights
/// are all zero.
///
/// # Variants
///
/// * `UseLocalMean` - Compute simple average of y-values in window (default)
/// * `ReturnOriginal` - Return the original unsmoothed y-value
/// * `ReturnNone` - Signal fitting failure by returning None
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum ZeroWeightFallback {
    /// Use local mean (default).
    ///
    /// This is the default and most robust option. When all weights are zero,
    /// compute the simple average of y-values in the window. This provides a
    /// reasonable fallback that smooths toward local data.
    #[default]
    UseLocalMean,

    /// Return the original y-value.
    ///
    /// When all weights are zero, return the original unsmoothed y-value
    /// for this point. This preserves the original data when smoothing fails.
    ReturnOriginal,

    /// Return None (propagate failure).
    ///
    /// When all weights are zero, return None to signal that fitting failed.
    /// The caller must handle this case explicitly, typically by deciding
    /// whether to skip the point or use a fallback strategy.
    ReturnNone,
}

impl ZeroWeightFallback {
    /// Create from u8 flag for backward compatibility.
    ///
    /// Allows interfacing with code that uses numeric flags:
    /// * 0 → UseLocalMean
    /// * 1 → ReturnOriginal
    /// * 2 → ReturnNone
    /// * other → UseLocalMean (default)
    #[inline]
    pub fn from_u8(flag: u8) -> Self {
        match flag {
            0 => ZeroWeightFallback::UseLocalMean,
            1 => ZeroWeightFallback::ReturnOriginal,
            2 => ZeroWeightFallback::ReturnNone,
            _ => ZeroWeightFallback::UseLocalMean, // Default for unknown values
        }
    }

    /// Convert to u8 flag for backward compatibility.
    #[inline]
    pub fn to_u8(self) -> u8 {
        match self {
            ZeroWeightFallback::UseLocalMean => 0,
            ZeroWeightFallback::ReturnOriginal => 1,
            ZeroWeightFallback::ReturnNone => 2,
        }
    }
}

// ============================================================================
// Weight Parameters
// ============================================================================

/// Parameters for weight computation.
///
/// This struct encapsulates the parameters needed to compute kernel weights
/// for a single point in the LOWESS smoothing process.
///
/// # Implementation Note
///
/// The thresholds `h1` and `h9` are implementation optimizations:
/// * `h1 = 0.001 × window_radius`: Points closer than this get weight 1.0
///   (avoids near-zero distance numerical issues)
/// * `h9 = 0.999 × window_radius`: Points farther than this get weight 0.0
///   (early termination for efficiency in sorted arrays)
///
/// These are not formal LOWESS parameters but implementation details for
/// numerical stability and performance.
pub struct WeightParams<T: Float> {
    /// Current x-value being fitted
    pub x_current: T,

    /// Window radius - defines the scale of the local fit
    pub window_radius: T,

    /// Near-zero threshold (0.001 × radius) for assigning weight 1.0
    pub h1: T,

    /// Far threshold (0.999 × radius) for early termination
    pub h9: T,

    /// Whether to use robustness weights
    pub use_robustness: bool,
}

impl<T: Float> WeightParams<T> {
    /// Construct WeightParams with validated window radius.
    ///
    /// # Parameters
    ///
    /// * `x_current` - The x-value of the point being fitted
    /// * `window_radius` - Radius for the local fit
    /// * `use_robustness` - Whether to incorporate robustness weights
    ///
    /// # Panics
    ///
    /// Debug builds assert that radius > 0. Release builds clamp tiny
    /// or zero radius to a small epsilon for numerical stability.
    pub fn new(x_current: T, window_radius: T, use_robustness: bool) -> Self {
        debug_assert!(
            window_radius > T::zero(),
            "WeightParams::new: window_radius must be positive"
        );

        // In release builds avoid panic: clamp tiny/zero radius to small epsilon
        let radius = if window_radius > T::zero() {
            window_radius
        } else {
            // Small absolute fallback for numerical stability
            T::from(1e-12).unwrap()
        };

        let h1 = T::from(0.001).unwrap() * radius;
        let h9 = T::from(0.999).unwrap() * radius;

        Self {
            x_current,
            window_radius: radius,
            h1,
            h9,
            use_robustness,
        }
    }
}

// ============================================================================
// Generalized Least Squares Model
// ============================================================================

/// Generalized Least Squares (GLS) model.
///
/// Encapsulates the result of a linear regression fit (slope and intercept),
/// along with the centroids (weighted means) of the fitted data.
///
/// This struct consolidates Ordinary Least Squares (OLS) and Weighted Least
/// Squares (WLS) results, providing a unified interface for prediction.
///
/// # Fields
///
/// * `slope` - β₁ (slope coefficient)
/// * `intercept` - β₀ (intercept coefficient)
/// * `x_mean` - Weighted mean of x-values
/// * `y_mean` - Weighted mean of y-values
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct GLSModel<T: Float> {
    /// Slope (β₁)
    pub slope: T,

    /// Intercept (β₀)
    pub intercept: T,

    /// Weighted mean of x-values
    pub x_mean: T,

    /// Weighted mean of y-values
    pub y_mean: T,
}

impl<T: Float> GLSModel<T> {
    /// Predict y-value for a given x.
    ///
    /// # Parameters
    ///
    /// * `x` - Independent variable value
    ///
    /// # Returns
    ///
    /// Predicted y-value: ŷ = β₀ + β₁x
    fn predict(&self, x: T) -> T {
        self.intercept + self.slope * x
    }

    /// Fit Ordinary Least Squares (OLS) regression.
    ///
    /// Fits the model y = β₀ + β₁x using unweighted least squares.
    ///
    /// # Parameters
    ///
    /// * `x` - Independent variable values
    /// * `y` - Dependent variable values
    ///
    /// # Returns
    ///
    /// [`GLSModel`] containing slope and intercept, or default values if fit
    /// fails (e.g., variance is 0).
    ///
    /// # Algorithm
    ///
    /// 1. Compute means: x̄ = mean(x), ȳ = mean(y)
    /// 2. Compute variance: Var(x) = Σ(x_i - x̄)²
    /// 3. Compute covariance: Cov(x,y) = Σ(x_i - x̄)(y_i - ȳ)
    /// 4. Compute slope: β₁ = Cov(x,y) / Var(x)
    /// 5. Compute intercept: β₀ = ȳ - β₁x̄
    fn fit_ols(x: &[T], y: &[T]) -> Self
    where
        T: fmt::Debug,
    {
        let n = x.len();
        if n == 0 {
            return GLSModel {
                slope: T::zero(),
                intercept: T::zero(),
                x_mean: T::zero(),
                y_mean: T::zero(),
            };
        }

        let n_t = T::from(n).unwrap_or(T::one());

        let mut sum_x = T::zero();
        let mut sum_y = T::zero();

        for i in 0..n {
            sum_x = sum_x + x[i];
            sum_y = sum_y + y[i];
        }

        let x_mean = sum_x / n_t;
        let y_mean = sum_y / n_t;

        let mut variance = T::zero();
        let mut covariance = T::zero();

        for i in 0..n {
            let dx = x[i] - x_mean;
            let dy = y[i] - y_mean;
            variance = variance + dx * dx;
            covariance = covariance + dx * dy;
        }

        let tol = T::from(1e-12).unwrap();
        if variance <= tol {
            // Fallback: horizontal line at y_mean
            return GLSModel {
                slope: T::zero(),
                intercept: y_mean,
                x_mean,
                y_mean,
            };
        }

        let slope = covariance / variance;
        let intercept = y_mean - slope * x_mean;

        GLSModel {
            slope,
            intercept,
            x_mean,
            y_mean,
        }
    }

    /// Compute global OLS linear fit for the full dataset.
    ///
    /// # Parameters
    ///
    /// * `x` - Independent variable values
    /// * `y` - Dependent variable values
    ///
    /// # Returns
    ///
    /// Vector of fitted y-values for each input x.
    pub fn global_ols(x: &[T], y: &[T]) -> Vec<T>
    where
        T: fmt::Debug,
    {
        let model = Self::fit_ols(x, y);
        x.iter().map(|&xi| model.predict(xi)).collect()
    }

    /// Fit Weighted Least Squares (WLS) regression.
    ///
    /// Fits the model y = β₀ + β₁x using weighted least squares with
    /// numerical stability checks.
    ///
    /// # Parameters
    ///
    /// * `x` - Independent variable values
    /// * `y` - Dependent variable values
    /// * `weights` - Weights for each point (should sum to 1)
    /// * `window_radius` - Window radius for numerical stability check
    ///
    /// # Returns
    ///
    /// [`GLSModel`] containing slope and intercept, or fallback values if fit
    /// fails.
    ///
    /// # Algorithm
    ///
    /// 1. Compute weighted means: x̄ = Σw_i x_i, ȳ = Σw_i y_i
    /// 2. Compute weighted variance: Var_w(x) = Σw_i(x_i - x̄)²
    /// 3. Compute weighted covariance: Cov_w(x,y) = Σw_i(x_i - x̄)(y_i - ȳ)
    /// 4. Compute slope: β₁ = Cov_w(x,y) / Var_w(x)
    /// 5. Compute intercept: β₀ = ȳ - β₁x̄
    fn fit_wls(x: &[T], y: &[T], weights: &[T], window_radius: T) -> GLSModel<T> {
        let n = x.len();
        if n == 0 {
            return GLSModel {
                slope: T::zero(),
                intercept: T::zero(),
                x_mean: T::zero(),
                y_mean: T::zero(),
            };
        }

        // Compute weighted means (assumes weights sum to 1)
        let mut x_mean = T::zero();
        let mut y_mean = T::zero();
        for i in 0..n {
            let w = weights[i];
            x_mean = x_mean + w * x[i];
            y_mean = y_mean + w * y[i];
        }

        // Degenerate window_radius: simple weighted average
        if window_radius <= T::zero() {
            return GLSModel {
                slope: T::zero(),
                intercept: y_mean,
                x_mean,
                y_mean,
            };
        }

        // Compute weighted covariance and variance
        let mut variance = T::zero(); // Σw(x - x̄)²
        let mut covariance = T::zero(); // Σw(x - x̄)(y - ȳ)

        for i in 0..n {
            let w = weights[i];
            let dx = x[i] - x_mean;
            let dy = y[i] - y_mean;
            variance = variance + w * dx * dx;
            covariance = covariance + w * dx * dy;
        }

        // Numerical stability check for variance
        let abs_tol = T::from(1e-7).unwrap();
        let rel_tol = T::epsilon() * window_radius * window_radius;
        let tol = abs_tol.max(rel_tol);

        if variance <= tol {
            return GLSModel {
                slope: T::zero(),
                intercept: y_mean,
                x_mean,
                y_mean,
            };
        }

        let slope = covariance / variance;
        let intercept = y_mean - slope * x_mean;

        GLSModel {
            slope,
            intercept,
            x_mean,
            y_mean,
        }
    }

    /// Weighted linear regression evaluated at a specific point.
    ///
    /// Fits a local weighted line ŷ = β₀ + β₁x to points in [left, right]
    /// and evaluates it at x_current.
    ///
    /// # Parameters
    ///
    /// * `x` - Independent variable values
    /// * `y` - Dependent variable values
    /// * `weights` - Normalized weights (should sum to 1 over [left, right])
    /// * `left` - Start index of window
    /// * `right` - End index of window
    /// * `x_current` - Point at which to evaluate the fit
    /// * `window_radius` - Window radius (for numerical stability checks)
    ///
    /// # Returns
    ///
    /// Fitted value at x_current.
    ///
    /// # Algorithm
    ///
    /// Uses [`GLSModel::fit_wls`] on the window slice [left, right].
    pub fn local_wls(
        x: &[T],
        y: &[T],
        weights: &[T],
        left: usize,
        right: usize,
        x_current: T,
        window_radius: T,
    ) -> T {
        let window_x = &x[left..=right];
        let window_y = &y[left..=right];
        let window_weights = &weights[left..=right];

        let model = Self::fit_wls(window_x, window_y, window_weights, window_radius);
        model.predict(x_current)
    }
}
