//! Cross-validation for LOWESS bandwidth selection.
//!
//! ## Purpose
//!
//! This module provides cross-validation tools for selecting the optimal
//! smoothing fraction (bandwidth) in LOWESS regression. It implements
//! generic k-fold and leave-one-out cross-validation strategies that work
//! with any smoothing function.
//!
//! ## Design notes
//!
//! * Cross-validation is generic over the smoothing implementation via callbacks.
//! * Supports both k-fold CV and leave-one-out CV (LOOCV).
//! * Uses linear interpolation to predict on held-out test data.
//! * Selects the fraction that minimizes root mean squared error (RMSE).
//! * All operations are generic over `Float` types to support f32 and f64.
//! * Supports both `std` and `no_std` environments.
//!
//! ## Available methods
//!
//! * **K-Fold CV**: Partitions data into k subsamples, training on k-1 and validating on 1
//! * **Leave-One-Out CV (LOOCV)**: Extreme case of k-fold where k equals sample size
//!
//! ## Key concepts
//!
//! ### Cross-Validation
//! Cross-validation estimates prediction error by repeatedly fitting the model
//! on training subsets and evaluating on held-out test subsets. The fraction
//! with the lowest average error is selected.
//!
//! ### K-Fold Strategy
//! Data is partitioned into k equal-sized folds. Each fold serves as the test
//! set once while the remaining k-1 folds form the training set. This provides
//! a good balance between computational cost and accuracy.
//!
//! ### Leave-One-Out Strategy
//! Each data point is held out once as the test set while all other points
//! form the training set. This is the most accurate but also the most
//! computationally expensive method (n iterations).
//!
//! ### Interpolation for Prediction
//! Since test points may not be in the training set, predictions are made
//! by linearly interpolating the smoothed training values. Constant
//! extrapolation is used at boundaries.
//!
//! ## Invariants
//!
//! * Training and test sets are disjoint in each fold.
//! * All data points are used for testing exactly once.
//! * The best fraction minimizes RMSE across all folds.
//! * Interpolation is well-defined for sorted x values.
//!
//! ## Non-goals
//!
//! * This module does not perform the actual smoothing (done via callback).
//! * This module does not validate input data or fraction ranges.
//! * This module does not support stratified or time-series CV.
//! * This module does not provide confidence intervals for CV scores.
//!
//! ## Visibility
//!
//! The [`CVMethod`] enum is part of the public API, allowing users to
//! configure cross-validation strategies. The internal implementation
//! details may change without notice.

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

#[cfg(feature = "std")]
use std::vec::Vec;

use core::cmp::Ordering;
use num_traits::Float;

// ============================================================================
// Cross-Validation Configuration
// ============================================================================

/// Cross-validation method configuration.
///
/// Specifies the cross-validation strategy to use for bandwidth selection.
/// Each variant defines how the data is partitioned into training and test sets.
///
/// # Variants
///
/// * `KFold(k)` - K-fold cross-validation with k folds
/// * `LOOCV` - Leave-one-out cross-validation (equivalent to n-fold)
///
/// # Choosing a Method
///
/// * **K-Fold**: Good balance between accuracy and speed. Common choices:
///   - k=5: Fast, reasonable accuracy
///   - k=10: Standard choice, good accuracy
///   - k=20: Higher accuracy, slower
///
/// * **LOOCV**: Maximum accuracy but computationally expensive (O(n²) evaluations).
///   Best for small datasets (n < 100) where accuracy is critical.
#[allow(clippy::upper_case_acronyms)]
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum CVMethod {
    /// K-fold cross-validation with k folds.
    ///
    /// Data is partitioned into k equal-sized subsets. Each subset serves
    /// as the test set once while the remaining k-1 subsets form the training set.
    ///
    /// Computational cost: O(k × m) where m is the number of candidate fractions.
    KFold(usize),

    /// Leave-one-out cross-validation.
    ///
    /// Each data point is held out once as the test set while all other
    /// points form the training set. This is equivalent to k-fold with k=n.
    ///
    /// Computational cost: O(n × m) where n is the dataset size and m is
    /// the number of candidate fractions.
    LOOCV,
}

impl CVMethod {
    // ========================================================================
    // Public API
    // ========================================================================

    /// Run cross-validation to select the best fraction.
    ///
    /// Executes the configured cross-validation strategy and returns the
    /// fraction that minimizes root mean squared error (RMSE) on held-out
    /// test data.
    ///
    /// # Parameters
    ///
    /// * `x` - Sorted x-coordinates (independent variable)
    /// * `y` - Y-coordinates (dependent variable)
    /// * `fractions` - Candidate smoothing fractions to evaluate
    /// * `smoother` - Closure that performs smoothing: `(x, y, frac) -> smoothed_y`
    ///
    /// # Returns
    ///
    /// Tuple of `(best_fraction, cv_scores)`:
    /// - `best_fraction`: The fraction with the lowest RMSE
    /// - `cv_scores`: RMSE for each candidate fraction (same order as input)
    ///
    /// # Algorithm
    ///
    /// 1. For each candidate fraction:
    ///    - Partition data into training and test sets
    ///    - Fit smoother on training data
    ///    - Predict on test data using interpolation
    ///    - Compute RMSE on test predictions
    /// 2. Select the fraction with the minimum RMSE
    pub fn run<T, F>(self, x: &[T], y: &[T], fractions: &[T], smoother: F) -> (T, Vec<T>)
    where
        T: Float,
        F: Fn(&[T], &[T], T) -> Vec<T> + Copy,
    {
        match self {
            CVMethod::KFold(k) => Self::kfold_cross_validation(x, y, fractions, k, smoother),
            CVMethod::LOOCV => Self::leave_one_out_cross_validation(x, y, fractions, smoother),
        }
    }

    // ========================================================================
    // Utility Methods
    // ========================================================================

    /// Build a data subset from a list of indices.
    ///
    /// Creates new vectors containing only the elements at the specified indices.
    /// Used to construct training and test sets for cross-validation.
    ///
    /// # Parameters
    ///
    /// * `x` - Full x-coordinate array
    /// * `y` - Full y-coordinate array
    /// * `indices` - Indices to extract
    ///
    /// # Returns
    ///
    /// Tuple of `(subset_x, subset_y)` containing the selected elements.
    pub fn build_subset_from_indices<T: Float>(
        x: &[T],
        y: &[T],
        indices: &[usize],
    ) -> (Vec<T>, Vec<T>) {
        let mut tx = Vec::with_capacity(indices.len());
        let mut ty = Vec::with_capacity(indices.len());
        for &i in indices {
            tx.push(x[i]);
            ty.push(y[i]);
        }
        (tx, ty)
    }

    /// Interpolate prediction at a new x value given fitted training values.
    ///
    /// Used for predicting on held-out test data in cross-validation.
    /// Performs linear interpolation between bracketing points or constant
    /// extrapolation at boundaries.
    ///
    /// # Parameters
    ///
    /// * `x_train` - Training x-coordinates (should be sorted)
    /// * `y_train` - Fitted training y-values (smoothed)
    /// * `x_new` - New x-coordinate to predict at
    ///
    /// # Returns
    ///
    /// Predicted y-value at `x_new`.
    ///
    /// # Algorithm
    ///
    /// 1. Handle edge cases (empty, single point)
    /// 2. Use constant extrapolation at boundaries
    /// 3. Binary search to find bracketing points
    /// 4. Linear interpolation: `y = y0 + α(y1 - y0)` where `α = (x - x0)/(x1 - x0)`
    ///
    /// # Implementation notes
    ///
    /// * Uses binary search for O(log n) bracketing
    /// * Handles duplicate x-values by averaging y-values
    /// * Constant extrapolation prevents unbounded predictions
    pub fn interpolate_prediction<T: Float>(x_train: &[T], y_train: &[T], x_new: T) -> T {
        let n = x_train.len();

        // Edge case: empty training set
        if n == 0 {
            return T::zero();
        }

        // Edge case: single training point
        if n == 1 {
            return y_train[0];
        }

        // Boundary handling: constant extrapolation
        if x_new <= x_train[0] {
            return y_train[0];
        }
        if x_new >= x_train[n - 1] {
            return y_train[n - 1];
        }

        // Binary search for bracketing points
        let mut left = 0;
        let mut right = n - 1;

        while right - left > 1 {
            let mid = (left + right) / 2;
            if x_train[mid] <= x_new {
                left = mid;
            } else {
                right = mid;
            }
        }

        // Linear interpolation between left and right
        let x0 = x_train[left];
        let x1 = x_train[right];
        let y0 = y_train[left];
        let y1 = y_train[right];

        let denom = x1 - x0;
        if denom <= T::zero() {
            // Duplicate x-values: return average of y-values
            return (y0 + y1) / T::from(2.0).unwrap();
        }

        let alpha = (x_new - x0) / denom;
        y0 + alpha * (y1 - y0)
    }

    // ========================================================================
    // Internal Cross-Validation Implementations
    // ========================================================================

    /// Select the best fraction based on cross-validation scores.
    ///
    /// Finds the fraction with the lowest RMSE score.
    ///
    /// # Parameters
    ///
    /// * `fractions` - Candidate fractions
    /// * `scores` - RMSE scores for each fraction
    ///
    /// # Returns
    ///
    /// Tuple of `(best_fraction, all_scores)`.
    fn select_best_fraction<T: Float>(fractions: &[T], scores: &[T]) -> (T, Vec<T>) {
        if fractions.is_empty() {
            return (T::zero(), Vec::new());
        }

        let best_idx = scores
            .iter()
            .enumerate()
            .min_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(Ordering::Equal))
            .map(|(i, _)| i)
            .unwrap_or(0);

        (fractions[best_idx], scores.to_vec())
    }

    /// Perform k-fold cross-validation.
    ///
    /// # Parameters
    ///
    /// * `x` - Sorted x-coordinates
    /// * `y` - Y-coordinates
    /// * `fractions` - Candidate fractions to evaluate
    /// * `k` - Number of folds
    /// * `smoother` - Smoothing function
    ///
    /// # Returns
    ///
    /// Tuple of `(best_fraction, cv_scores)`.
    ///
    /// # Algorithm
    ///
    /// For each fraction:
    /// 1. Partition data into k folds
    /// 2. For each fold:
    ///    - Use fold as test set, remaining k-1 folds as training set
    ///    - Fit smoother on training data
    ///    - Predict on test data via interpolation
    ///    - Accumulate squared errors
    /// 3. Compute RMSE = sqrt(mean(squared_errors))
    fn kfold_cross_validation<T, F>(
        x: &[T],
        y: &[T],
        fractions: &[T],
        k: usize,
        smoother: F,
    ) -> (T, Vec<T>)
    where
        T: Float,
        F: Fn(&[T], &[T], T) -> Vec<T>,
    {
        if fractions.is_empty() || k < 2 {
            let default = T::from(0.67).unwrap_or(T::from(0.5).unwrap());
            return (default, vec![T::zero()]);
        }

        let n = x.len();
        let fold_size = n / k;

        let mut cv_scores = vec![T::zero(); fractions.len()];

        for (frac_idx, &frac) in fractions.iter().enumerate() {
            let mut total_error = T::zero();
            let mut count = T::zero();

            for fold in 0..k {
                // Define test set for this fold
                let test_start = fold * fold_size;
                let test_end = if fold == k - 1 {
                    n // Last fold includes remainder
                } else {
                    (fold + 1) * fold_size
                };

                // Build training set (all indices except test set)
                let train_indices: Vec<usize> = (0..n)
                    .filter(|&i| i < test_start || i >= test_end)
                    .collect();
                let (train_x, train_y) = Self::build_subset_from_indices(x, y, &train_indices);

                // Fit smoother on training data
                let train_smooth = smoother(&train_x, &train_y, frac);

                // Compute error on test set using interpolation
                for i in test_start..test_end {
                    let predicted = Self::interpolate_prediction(&train_x, &train_smooth, x[i]);
                    let error = y[i] - predicted;
                    total_error = total_error + error * error;
                    count = count + T::one();
                }
            }

            // Compute RMSE for this fraction
            if count > T::zero() {
                cv_scores[frac_idx] = (total_error / count).sqrt();
            } else {
                cv_scores[frac_idx] = T::infinity();
            }
        }

        Self::select_best_fraction(fractions, &cv_scores)
    }

    /// Perform leave-one-out cross-validation (LOOCV).
    ///
    /// # Parameters
    ///
    /// * `x` - Sorted x-coordinates
    /// * `y` - Y-coordinates
    /// * `fractions` - Candidate fractions to evaluate
    /// * `smoother` - Smoothing function
    ///
    /// # Returns
    ///
    /// Tuple of `(best_fraction, cv_scores)`.
    ///
    /// # Algorithm
    ///
    /// For each fraction:
    /// 1. For each data point i:
    ///    - Use point i as test set, all other points as training set
    ///    - Fit smoother on training data
    ///    - Predict at x[i] via interpolation
    ///    - Accumulate squared error
    /// 2. Compute RMSE = sqrt(mean(squared_errors))
    fn leave_one_out_cross_validation<T, F>(
        x: &[T],
        y: &[T],
        fractions: &[T],
        smoother: F,
    ) -> (T, Vec<T>)
    where
        T: Float,
        F: Fn(&[T], &[T], T) -> Vec<T>,
    {
        if fractions.is_empty() {
            let default = T::from(0.67).unwrap_or(T::from(0.5).unwrap());
            return (default, vec![T::zero()]);
        }

        let n = x.len();
        if n < 2 {
            return (fractions[0], vec![T::zero(); fractions.len()]);
        }

        let mut cv_scores = vec![T::zero(); fractions.len()];

        for (frac_idx, &frac) in fractions.iter().enumerate() {
            let mut total_error = T::zero();

            for i in 0..n {
                // Build leave-one-out training set (all points except i)
                let train_indices: Vec<usize> = (0..n).filter(|&j| j != i).collect();
                let (train_x, train_y) = Self::build_subset_from_indices(x, y, &train_indices);

                // Fit smoother on training data
                let train_smooth = smoother(&train_x, &train_y, frac);

                // Predict at held-out point
                let predicted = Self::interpolate_prediction(&train_x, &train_smooth, x[i]);
                let error = y[i] - predicted;
                total_error = total_error + error * error;
            }

            // Compute RMSE for this fraction
            cv_scores[frac_idx] = (total_error / T::from(n).unwrap()).sqrt();
        }

        Self::select_best_fraction(fractions, &cv_scores)
    }
}
