//! Diagnostic metrics for LOWESS fit quality assessment.
//!
//! ## Purpose
//!
//! This module provides comprehensive diagnostic tools for evaluating the
//! quality of LOWESS smoothing results. It computes goodness-of-fit metrics,
//! model selection criteria, and coverage statistics to help assess and
//! compare different smoothing configurations.
//!
//! ## Design notes
//!
//! * All metrics are computed from residuals (y - ŷ) and fitted values.
//! * Uses MAD (Median Absolute Deviation) for robust scale estimation.
//! * Effective degrees of freedom are estimated from standard errors.
//! * AIC and AICc provide model selection criteria for bandwidth tuning.
//! * All computations are generic over `Float` types to support f32 and f64.
//!
//! ## Available metrics
//!
//! * **Error metrics**: RMSE, MAE, RSS
//! * **Goodness-of-fit**: R² (coefficient of determination)
//! * **Model selection**: AIC, AICc (corrected AIC)
//! * **Robustness**: Residual standard deviation (from MAD)
//! * **Complexity**: Effective degrees of freedom
//! * **Coverage**: Interval coverage probability
//!
//! ## Key concepts
//!
//! ### Residual Metrics
//! Residual-based metrics (RMSE, MAE, RSS) measure the magnitude of
//! prediction errors. Lower values indicate better fit.
//!
//! ### R-Squared
//! The coefficient of determination measures the proportion of variance
//! explained by the model. Values range from 0 (no fit) to 1 (perfect fit).
//!
//! ### Effective Degrees of Freedom
//! Estimated from the trace of the smoother matrix (sum of leverages).
//! Represents the model complexity, accounting for the adaptive nature
//! of LOWESS.
//!
//! ### AIC and AICc
//! Akaike Information Criterion balances fit quality against model complexity.
//! AICc includes a finite-sample correction. Lower values indicate better
//! models for the given data.
//!
//! ### MAD-based Standard Deviation
//! Uses MAD × 1.4826 to estimate residual standard deviation robustly,
//! resistant to outliers.
//!
//! ## Invariants
//!
//! * All metrics are non-negative (except R² which can be negative for poor fits).
//! * RMSE ≥ 0, MAE ≥ 0, RSS ≥ 0.
//! * 0 ≤ coverage ≤ 1.
//! * Effective DF > 0 when computable.
//!
//! ## Non-goals
//!
//! * This module does not perform the smoothing itself.
//! * This module does not validate input data.
//! * This module does not provide hypothesis tests or p-values.
//! * This module does not compute weighted metrics.
//!
//! ## Visibility
//!
//! The [`Diagnostics`] struct is part of the public API, allowing users
//! to assess fit quality. Individual metric functions are also public
//! for standalone use.

use crate::math::mad::compute_mad;
use core::fmt;
use num_traits::Float;

// ============================================================================
// Diagnostics Structure
// ============================================================================

/// Diagnostic statistics for LOWESS fit quality.
///
/// Contains comprehensive metrics for assessing the quality of a LOWESS
/// smoothing result, including error measures, goodness-of-fit, and
/// model selection criteria.
///
/// # Fields
///
/// * `rmse` - Root Mean Squared Error
/// * `mae` - Mean Absolute Error
/// * `r_squared` - Coefficient of determination (R²)
/// * `aic` - Akaike Information Criterion (if effective DF available)
/// * `aicc` - Corrected AIC (if effective DF available)
/// * `effective_df` - Effective degrees of freedom (if computable)
/// * `residual_sd` - Residual standard deviation (from MAD)
#[derive(Debug, Clone, PartialEq)]
pub struct Diagnostics<T> {
    /// Root Mean Squared Error: sqrt(mean(residuals²))
    pub rmse: T,

    /// Mean Absolute Error: mean(|residuals|)
    pub mae: T,

    /// R-squared (coefficient of determination): 1 - SS_res/SS_tot
    pub r_squared: T,

    /// Akaike Information Criterion (requires effective DF)
    pub aic: Option<T>,

    /// Corrected AIC with finite-sample correction (requires effective DF)
    pub aicc: Option<T>,

    /// Effective degrees of freedom (estimated from leverages)
    pub effective_df: Option<T>,

    /// Residual standard deviation (estimated from MAD)
    pub residual_sd: T,
}

impl<T: Float> Diagnostics<T> {
    // ========================================================================
    // Constants
    // ========================================================================

    /// Number of parameters in local linear regression (intercept + slope).
    const LINEAR_PARAMS: f64 = 2.0;

    /// Minimum tuned-scale absolute epsilon to avoid division-by-zero.
    const MIN_TUNED_SCALE: f64 = 1e-12;

    /// Constant to convert MAD to an unbiased estimate of σ for normal data.
    ///
    /// For normally distributed data, MAD × 1.4826 ≈ standard deviation.
    const MAD_TO_STD_FACTOR: f64 = 1.4826;

    // ========================================================================
    // Main Computation
    // ========================================================================

    /// Compute diagnostic statistics from fit results.
    ///
    /// Computes all available diagnostic metrics from the smoothing results.
    /// AIC and AICc are only computed if standard errors are provided
    /// (allowing effective DF estimation).
    ///
    /// # Parameters
    ///
    /// * `y` - Original y-values
    /// * `y_smooth` - Smoothed y-values
    /// * `residuals` - Residuals (y - y_smooth)
    /// * `std_errors` - Standard errors of the fit (optional, for effective DF)
    ///
    /// # Returns
    ///
    /// [`Diagnostics`] struct populated with computed metrics.
    ///
    /// # Algorithm
    ///
    /// 1. Compute error metrics (RMSE, MAE)
    /// 2. Compute R²
    /// 3. Estimate residual SD from MAD
    /// 4. If standard errors provided:
    ///    - Estimate effective DF
    ///    - Compute AIC and AICc
    pub fn compute(y: &[T], y_smooth: &[T], residuals: &[T], std_errors: Option<&[T]>) -> Self {
        let rmse = Self::calculate_rmse(y, y_smooth);
        let mae = Self::calculate_mae(y, y_smooth);
        let r_squared = Self::calculate_r_squared(y, y_smooth);
        let residual_sd = Self::calculate_residual_sd(residuals);

        let effective_df = std_errors.and_then(|se| Self::calculate_effective_df(se, residual_sd));

        let (aic, aicc) = if let Some(df) = effective_df {
            (
                Some(Self::calculate_aic(residuals, df)),
                Some(Self::calculate_aicc(residuals, df)),
            )
        } else {
            (None, None)
        };

        Diagnostics {
            rmse,
            mae,
            r_squared,
            aic,
            aicc,
            effective_df,
            residual_sd,
        }
    }

    // ========================================================================
    // Error Metrics
    // ========================================================================

    /// Compute residual sum of squares (RSS).
    ///
    /// RSS = Σ(residuals²)
    ///
    /// # Parameters
    ///
    /// * `residuals` - Residual values (y - ŷ)
    ///
    /// # Returns
    ///
    /// Sum of squared residuals.
    fn calculate_rss(residuals: &[T]) -> T {
        residuals.iter().fold(T::zero(), |acc, &r| acc + r * r)
    }

    /// Compute Root Mean Squared Error (RMSE).
    ///
    /// RMSE = sqrt(mean((y - ŷ)²))
    ///
    /// # Parameters
    ///
    /// * `y` - Original y-values
    /// * `y_smooth` - Smoothed y-values
    ///
    /// # Returns
    ///
    /// Root mean squared error.
    pub fn calculate_rmse(y: &[T], y_smooth: &[T]) -> T {
        let n = y.len();
        if n == 0 {
            return T::zero();
        }

        let n_t = T::from(n).unwrap_or(T::one());
        let rss = y
            .iter()
            .zip(y_smooth.iter())
            .fold(T::zero(), |acc, (&yi, &ys)| {
                let r = yi - ys;
                acc + r * r
            });

        (rss / n_t).sqrt()
    }

    /// Compute Mean Absolute Error (MAE).
    ///
    /// MAE = mean(|y - ŷ|)
    ///
    /// # Parameters
    ///
    /// * `y` - Original y-values
    /// * `y_smooth` - Smoothed y-values
    ///
    /// # Returns
    ///
    /// Mean absolute error.
    pub fn calculate_mae(y: &[T], y_smooth: &[T]) -> T {
        let n = y.len();
        if n == 0 {
            return T::zero();
        }

        let n_t = T::from(n).unwrap_or(T::one());
        let sum = y
            .iter()
            .zip(y_smooth.iter())
            .fold(T::zero(), |acc, (&yi, &ys)| acc + (yi - ys).abs());

        sum / n_t
    }

    // ========================================================================
    // Goodness-of-Fit Metrics
    // ========================================================================

    /// Compute unweighted R² (coefficient of determination).
    ///
    /// R² = 1 - SS_res/SS_tot
    ///
    /// where:
    /// - SS_res = Σ(y - ŷ)² (residual sum of squares)
    /// - SS_tot = Σ(y - ȳ)² (total sum of squares)
    ///
    /// # Parameters
    ///
    /// * `y` - Original y-values
    /// * `y_smooth` - Smoothed y-values
    ///
    /// # Returns
    ///
    /// R-squared value. Returns 1.0 for perfect fit, 0.0 for no improvement
    /// over mean, and can be negative for very poor fits.
    pub fn calculate_r_squared(y: &[T], y_smooth: &[T]) -> T {
        let n = y.len();
        if n == 0 {
            return T::zero();
        }
        if n == 1 {
            return T::one();
        }

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

        // Compute mean
        let sum = y.iter().copied().fold(T::zero(), |acc, v| acc + v);
        let mean = sum / n_t;

        // Compute SS_tot and SS_res in one pass
        let (ss_tot, ss_res) =
            y.iter()
                .zip(y_smooth.iter())
                .fold((T::zero(), T::zero()), |(tot, res), (&yi, &ys)| {
                    let deviation = yi - mean;
                    let residual = yi - ys;
                    (tot + deviation * deviation, res + residual * residual)
                });

        if ss_tot == T::zero() {
            // All y values are identical
            if ss_res == T::zero() {
                T::one() // Perfect fit
            } else {
                T::zero() // No variance to explain
            }
        } else {
            T::one() - ss_res / ss_tot
        }
    }

    // ========================================================================
    // Robust Scale Estimation
    // ========================================================================

    /// Estimate residual standard deviation using MAD.
    ///
    /// Uses MAD (Median Absolute Deviation) scaled by 1.4826 to provide
    /// a robust estimate of the residual standard deviation. This is
    /// resistant to outliers.
    ///
    /// # Parameters
    ///
    /// * `residuals` - Residual values (y - ŷ)
    ///
    /// # Returns
    ///
    /// Estimated residual standard deviation.
    ///
    /// # Algorithm
    ///
    /// 1. Compute MAD of residuals
    /// 2. Scale by 1.4826 (for normal data, MAD × 1.4826 ≈ σ)
    /// 3. Apply minimum threshold to avoid zero scale
    pub fn calculate_residual_sd(residuals: &[T]) -> T {
        let n = residuals.len();
        if n == 0 {
            return T::zero();
        }

        let scale_const = T::from(Self::MAD_TO_STD_FACTOR).unwrap();

        if n == 1 {
            return residuals[0].abs() * scale_const;
        }

        let mad = compute_mad(residuals);
        if mad > T::zero() {
            mad * scale_const
        } else {
            // Apply minimum scale to avoid division by zero
            let min_eps = T::from(Self::MIN_TUNED_SCALE).unwrap();
            min_eps * scale_const
        }
    }

    // ========================================================================
    // Model Complexity
    // ========================================================================

    /// Estimate effective degrees of freedom from standard errors.
    ///
    /// Estimates the effective DF by computing the sum of leverage values,
    /// where leverage is approximated from the ratio of standard errors
    /// to residual standard deviation.
    ///
    /// # Parameters
    ///
    /// * `std_errors` - Standard errors of the fit at each point
    /// * `residual_sd` - Residual standard deviation
    ///
    /// # Returns
    ///
    /// Effective degrees of freedom, or `None` if not computable.
    ///
    /// # Algorithm
    ///
    /// 1. For each point, compute leverage ≈ (SE / residual_sd)²
    /// 2. Clamp leverage to [0, 1]
    /// 3. Sum all leverages to get effective DF
    fn calculate_effective_df(std_errors: &[T], residual_sd: T) -> Option<T> {
        if residual_sd <= T::zero() {
            return None;
        }

        let leverage_sum = std_errors.iter().fold(T::zero(), |acc, &se| {
            let ratio = se / residual_sd;
            let mut leverage = ratio * ratio;

            // Clamp leverage to valid range [0, 1]
            if leverage < T::zero() {
                leverage = T::zero();
            } else if leverage > T::one() {
                leverage = T::one();
            }

            acc + leverage
        });

        if leverage_sum > T::zero() {
            Some(leverage_sum)
        } else {
            None
        }
    }

    /// Compute effective degrees of freedom from leverage values.
    ///
    /// Directly sums pre-computed leverage values (trace of smoother matrix).
    ///
    /// # Parameters
    ///
    /// * `leverage_values` - Leverage values for each point
    ///
    /// # Returns
    ///
    /// Sum of leverages (effective DF).
    pub fn calculate_effective_df_from_leverages(leverage_values: &[T]) -> T {
        leverage_values
            .iter()
            .copied()
            .fold(T::zero(), |acc, v| acc + v)
    }

    // ========================================================================
    // Model Selection Criteria
    // ========================================================================

    /// Compute Akaike Information Criterion (AIC).
    ///
    /// AIC = n × ln(RSS/n) + 2 × effective_df
    ///
    /// Lower AIC indicates a better model. Balances fit quality against
    /// model complexity.
    ///
    /// # Parameters
    ///
    /// * `residuals` - Residual values
    /// * `effective_df` - Effective degrees of freedom
    ///
    /// # Returns
    ///
    /// AIC value, or infinity if not computable.
    pub fn calculate_aic(residuals: &[T], effective_df: T) -> T {
        let n = T::from(residuals.len()).unwrap_or(T::one());
        let rss = Self::calculate_rss(residuals);

        if rss <= T::zero() || n <= T::zero() {
            return T::infinity();
        }

        n * (rss / n).ln() + T::from(Self::LINEAR_PARAMS).unwrap() * effective_df
    }

    /// Compute corrected AIC (AICc).
    ///
    /// AICc = AIC + (2 × k × (k + 1)) / (n - k - 1)
    ///
    /// where k = effective_df. Includes a finite-sample correction that
    /// is important for small sample sizes.
    ///
    /// # Parameters
    ///
    /// * `residuals` - Residual values
    /// * `effective_df` - Effective degrees of freedom
    ///
    /// # Returns
    ///
    /// AICc value, or infinity if not computable.
    pub fn calculate_aicc(residuals: &[T], effective_df: T) -> T {
        let n = T::from(residuals.len()).unwrap_or(T::one());
        let aic = Self::calculate_aic(residuals, effective_df);
        let denom = n - effective_df - T::one();

        if denom <= T::zero() {
            return T::infinity();
        }

        let correction =
            (T::from(Self::LINEAR_PARAMS).unwrap() * effective_df * (effective_df + T::one()))
                / denom;

        aic + correction
    }
}

// ============================================================================
// Display Implementation
// ============================================================================

impl<T: Float + fmt::Display> fmt::Display for Diagnostics<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        writeln!(f, "LOWESS Diagnostics:")?;
        writeln!(f, "  RMSE:         {:.6}", self.rmse)?;
        writeln!(f, "  MAE:          {:.6}", self.mae)?;
        writeln!(f, "  R²:           {:.6}", self.r_squared)?;
        writeln!(f, "  Residual SD:  {:.6}", self.residual_sd)?;

        if let Some(df) = self.effective_df {
            writeln!(f, "  Effective DF: {:.2}", df)?;
        }
        if let Some(aic) = self.aic {
            writeln!(f, "  AIC:          {:.2}", aic)?;
        }
        if let Some(aicc) = self.aicc {
            writeln!(f, "  AICc:         {:.2}", aicc)?;
        }

        Ok(())
    }
}
