//! Cross-validation framework for kernel parameter selection
//!
//! This module provides comprehensive cross-validation methods specifically designed
//! for kernel approximation methods and parameter selection.

use crate::{Nystroem, ParameterLearner, ParameterSet, RBFSampler};
use rayon::prelude::*;
use scirs2_core::ndarray::{s, Array1, Array2};
use scirs2_core::random::rngs::StdRng;
use scirs2_core::random::seq::SliceRandom;
use scirs2_core::random::{thread_rng, SeedableRng};
use sklears_core::{
    error::{Result, SklearsError},
    traits::{Fit, Transform},
};
use std::collections::HashMap;

/// Cross-validation strategies
#[derive(Debug, Clone)]
/// CVStrategy
pub enum CVStrategy {
    /// K-fold cross-validation
    KFold {
        /// Number of folds
        n_folds: usize,
        /// Shuffle data before folding
        shuffle: bool,
    },
    /// Stratified K-fold (for classification tasks)
    StratifiedKFold {
        /// Number of folds
        n_folds: usize,
        /// Shuffle data before folding
        shuffle: bool,
    },
    /// Leave-one-out cross-validation
    LeaveOneOut,
    /// Leave-P-out cross-validation
    LeavePOut {
        /// Number of points to leave out
        p: usize,
    },
    /// Time series split
    TimeSeriesSplit {
        /// Number of splits
        n_splits: usize,
        /// Maximum training size
        max_train_size: Option<usize>,
    },
    /// Monte Carlo cross-validation
    MonteCarlo {
        /// Number of random splits
        n_splits: usize,
        /// Test size fraction
        test_size: f64,
    },
}

/// Scoring metrics for cross-validation
#[derive(Debug, Clone)]
/// ScoringMetric
pub enum ScoringMetric {
    /// Kernel alignment score
    KernelAlignment,
    /// Mean squared error (for regression)
    MeanSquaredError,
    /// Mean absolute error (for regression)
    MeanAbsoluteError,
    /// R² score (for regression)
    R2Score,
    /// Accuracy (for classification)
    Accuracy,
    /// F1 score (for classification)
    F1Score,
    /// Log-likelihood
    LogLikelihood,
    /// Custom scoring function
    Custom,
}

/// Configuration for cross-validation
#[derive(Debug, Clone)]
/// CrossValidationConfig
pub struct CrossValidationConfig {
    /// Cross-validation strategy
    pub cv_strategy: CVStrategy,
    /// Scoring metric
    pub scoring_metric: ScoringMetric,
    /// Random seed for reproducibility
    pub random_seed: Option<u64>,
    /// Number of parallel jobs
    pub n_jobs: usize,
    /// Return training scores as well
    pub return_train_score: bool,
    /// Verbose output
    pub verbose: bool,
    /// Fit parameters for kernel methods
    pub fit_params: HashMap<String, f64>,
}

impl Default for CrossValidationConfig {
    fn default() -> Self {
        Self {
            cv_strategy: CVStrategy::KFold {
                n_folds: 5,
                shuffle: true,
            },
            scoring_metric: ScoringMetric::KernelAlignment,
            random_seed: None,
            n_jobs: num_cpus::get(),
            return_train_score: false,
            verbose: false,
            fit_params: HashMap::new(),
        }
    }
}

/// Results from cross-validation
#[derive(Debug, Clone)]
/// CrossValidationResult
pub struct CrossValidationResult {
    /// Test scores for each fold
    pub test_scores: Vec<f64>,
    /// Training scores for each fold (if requested)
    pub train_scores: Option<Vec<f64>>,
    /// Mean test score
    pub mean_test_score: f64,
    /// Standard deviation of test scores
    pub std_test_score: f64,
    /// Mean training score (if available)
    pub mean_train_score: Option<f64>,
    /// Standard deviation of training scores (if available)
    pub std_train_score: Option<f64>,
    /// Fit times for each fold
    pub fit_times: Vec<f64>,
    /// Score times for each fold
    pub score_times: Vec<f64>,
}

/// Cross-validation splitter interface
pub trait CVSplitter {
    /// Generate train/test indices for all splits
    fn split(&self, x: &Array2<f64>, y: Option<&Array1<f64>>) -> Vec<(Vec<usize>, Vec<usize>)>;
}

/// K-fold cross-validation splitter
pub struct KFoldSplitter {
    n_folds: usize,
    shuffle: bool,
    random_seed: Option<u64>,
}

impl KFoldSplitter {
    pub fn new(n_folds: usize, shuffle: bool, random_seed: Option<u64>) -> Self {
        Self {
            n_folds,
            shuffle,
            random_seed,
        }
    }
}

impl CVSplitter for KFoldSplitter {
    fn split(&self, x: &Array2<f64>, _y: Option<&Array1<f64>>) -> Vec<(Vec<usize>, Vec<usize>)> {
        let n_samples = x.nrows();
        let mut indices: Vec<usize> = (0..n_samples).collect();

        if self.shuffle {
            let mut rng = if let Some(seed) = self.random_seed {
                StdRng::seed_from_u64(seed)
            } else {
                StdRng::from_seed(thread_rng().random())
            };

            indices.shuffle(&mut rng);
        }

        let fold_size = n_samples / self.n_folds;
        let mut splits = Vec::new();

        for fold in 0..self.n_folds {
            let start = fold * fold_size;
            let end = if fold == self.n_folds - 1 {
                n_samples
            } else {
                (fold + 1) * fold_size
            };

            let test_indices = indices[start..end].to_vec();
            let train_indices = indices[..start]
                .iter()
                .chain(indices[end..].iter())
                .cloned()
                .collect();

            splits.push((train_indices, test_indices));
        }

        splits
    }
}

/// Time series cross-validation splitter
pub struct TimeSeriesSplitter {
    n_splits: usize,
    max_train_size: Option<usize>,
}

impl TimeSeriesSplitter {
    pub fn new(n_splits: usize, max_train_size: Option<usize>) -> Self {
        Self {
            n_splits,
            max_train_size,
        }
    }
}

impl CVSplitter for TimeSeriesSplitter {
    fn split(&self, x: &Array2<f64>, _y: Option<&Array1<f64>>) -> Vec<(Vec<usize>, Vec<usize>)> {
        let n_samples = x.nrows();
        let test_size = n_samples / (self.n_splits + 1);
        let mut splits = Vec::new();

        for split in 0..self.n_splits {
            let test_start = (split + 1) * test_size;
            let test_end = if split == self.n_splits - 1 {
                n_samples
            } else {
                (split + 2) * test_size
            };

            let train_end = test_start;
            let train_start = if let Some(max_size) = self.max_train_size {
                train_end.saturating_sub(max_size)
            } else {
                0
            };

            let train_indices = (train_start..train_end).collect();
            let test_indices = (test_start..test_end).collect();

            splits.push((train_indices, test_indices));
        }

        splits
    }
}

/// Monte Carlo cross-validation splitter
pub struct MonteCarloCVSplitter {
    n_splits: usize,
    test_size: f64,
    random_seed: Option<u64>,
}

impl MonteCarloCVSplitter {
    pub fn new(n_splits: usize, test_size: f64, random_seed: Option<u64>) -> Self {
        Self {
            n_splits,
            test_size,
            random_seed,
        }
    }
}

impl CVSplitter for MonteCarloCVSplitter {
    fn split(&self, x: &Array2<f64>, _y: Option<&Array1<f64>>) -> Vec<(Vec<usize>, Vec<usize>)> {
        let n_samples = x.nrows();
        let test_samples = (n_samples as f64 * self.test_size) as usize;
        let mut rng = if let Some(seed) = self.random_seed {
            StdRng::seed_from_u64(seed)
        } else {
            StdRng::from_seed(thread_rng().random())
        };

        let mut splits = Vec::new();

        for _ in 0..self.n_splits {
            let mut indices: Vec<usize> = (0..n_samples).collect();

            indices.shuffle(&mut rng);

            let test_indices = indices[..test_samples].to_vec();
            let train_indices = indices[test_samples..].to_vec();

            splits.push((train_indices, test_indices));
        }

        splits
    }
}

/// Main cross-validation framework
pub struct CrossValidator {
    config: CrossValidationConfig,
}

impl CrossValidator {
    /// Create a new cross-validator
    pub fn new(config: CrossValidationConfig) -> Self {
        Self { config }
    }

    /// Perform cross-validation for RBF sampler
    pub fn cross_validate_rbf(
        &self,
        x: &Array2<f64>,
        y: Option<&Array1<f64>>,
        parameters: &ParameterSet,
    ) -> Result<CrossValidationResult> {
        let splitter = self.create_splitter()?;
        let splits = splitter.split(x, y);

        if self.config.verbose {
            println!("Performing cross-validation with {} splits", splits.len());
        }

        // Parallel evaluation of each fold
        let fold_results: Result<Vec<_>> = splits
            .par_iter()
            .enumerate()
            .map(|(fold_idx, (train_indices, test_indices))| {
                let start_time = std::time::Instant::now();

                // Extract training and test data
                let x_train = self.extract_samples(x, train_indices);
                let x_test = self.extract_samples(x, test_indices);
                let y_train = y.map(|y_data| self.extract_targets(y_data, train_indices));
                let y_test = y.map(|y_data| self.extract_targets(y_data, test_indices));

                // Fit RBF sampler
                let sampler = RBFSampler::new(parameters.n_components).gamma(parameters.gamma);
                let fitted = sampler.fit(&x_train, &())?;
                let fit_time = start_time.elapsed().as_secs_f64();

                // Transform data
                let x_train_transformed = fitted.transform(&x_train)?;
                let x_test_transformed = fitted.transform(&x_test)?;

                // Compute scores
                let score_start = std::time::Instant::now();
                let test_score = self.compute_score(
                    &x_test,
                    &x_test_transformed,
                    y_test.as_ref(),
                    parameters.gamma,
                )?;

                let train_score = if self.config.return_train_score {
                    Some(self.compute_score(
                        &x_train,
                        &x_train_transformed,
                        y_train.as_ref(),
                        parameters.gamma,
                    )?)
                } else {
                    None
                };

                let score_time = score_start.elapsed().as_secs_f64();

                if self.config.verbose {
                    println!(
                        "Fold {}: test_score = {:.6}, fit_time = {:.3}s",
                        fold_idx, test_score, fit_time
                    );
                }

                Ok((test_score, train_score, fit_time, score_time))
            })
            .collect();

        let fold_results = fold_results?;

        self.aggregate_results(fold_results)
    }

    /// Perform cross-validation for Nyström method
    pub fn cross_validate_nystroem(
        &self,
        x: &Array2<f64>,
        y: Option<&Array1<f64>>,
        parameters: &ParameterSet,
    ) -> Result<CrossValidationResult> {
        use crate::nystroem::Kernel;

        let splitter = self.create_splitter()?;
        let splits = splitter.split(x, y);

        let fold_results: Result<Vec<_>> = splits
            .par_iter()
            .enumerate()
            .map(|(fold_idx, (train_indices, test_indices))| {
                let start_time = std::time::Instant::now();

                // Extract data
                let x_train = self.extract_samples(x, train_indices);
                let x_test = self.extract_samples(x, test_indices);
                let y_train = y.map(|y_data| self.extract_targets(y_data, train_indices));
                let y_test = y.map(|y_data| self.extract_targets(y_data, test_indices));

                // Fit Nyström
                let kernel = Kernel::Rbf {
                    gamma: parameters.gamma,
                };
                let nystroem = Nystroem::new(kernel, parameters.n_components);
                let fitted = nystroem.fit(&x_train, &())?;
                let fit_time = start_time.elapsed().as_secs_f64();

                // Transform data
                let x_train_transformed = fitted.transform(&x_train)?;
                let x_test_transformed = fitted.transform(&x_test)?;

                // Compute scores
                let score_start = std::time::Instant::now();
                let test_score = self.compute_score(
                    &x_test,
                    &x_test_transformed,
                    y_test.as_ref(),
                    parameters.gamma,
                )?;

                let train_score = if self.config.return_train_score {
                    Some(self.compute_score(
                        &x_train,
                        &x_train_transformed,
                        y_train.as_ref(),
                        parameters.gamma,
                    )?)
                } else {
                    None
                };

                let score_time = score_start.elapsed().as_secs_f64();

                if self.config.verbose {
                    println!("Fold {}: test_score = {:.6}", fold_idx, test_score);
                }

                Ok((test_score, train_score, fit_time, score_time))
            })
            .collect();

        let fold_results = fold_results?;
        self.aggregate_results(fold_results)
    }

    /// Cross-validate with parameter search
    pub fn cross_validate_with_search(
        &self,
        x: &Array2<f64>,
        y: Option<&Array1<f64>>,
        parameter_learner: &ParameterLearner,
    ) -> Result<(ParameterSet, CrossValidationResult)> {
        // First, optimize parameters using the parameter learner
        let optimization_result = parameter_learner.optimize_rbf_parameters(x, y)?;
        let best_params = optimization_result.best_parameters;

        if self.config.verbose {
            println!(
                "Best parameters found: gamma={:.6}, n_components={}",
                best_params.gamma, best_params.n_components
            );
        }

        // Then perform cross-validation with the best parameters
        let cv_result = self.cross_validate_rbf(x, y, &best_params)?;

        Ok((best_params, cv_result))
    }

    /// Grid search with cross-validation
    pub fn grid_search_cv(
        &self,
        x: &Array2<f64>,
        y: Option<&Array1<f64>>,
        param_grid: &HashMap<String, Vec<f64>>,
    ) -> Result<(
        ParameterSet,
        f64,
        HashMap<ParameterSet, CrossValidationResult>,
    )> {
        let gamma_values = param_grid.get("gamma").ok_or_else(|| {
            SklearsError::InvalidInput("gamma parameter missing from grid".to_string())
        })?;

        let n_components_values = param_grid
            .get("n_components")
            .ok_or_else(|| {
                SklearsError::InvalidInput("n_components parameter missing from grid".to_string())
            })?
            .iter()
            .map(|&x| x as usize)
            .collect::<Vec<_>>();

        let mut best_score = f64::NEG_INFINITY;
        let mut best_params = ParameterSet {
            gamma: gamma_values[0],
            n_components: n_components_values[0],
            degree: None,
            coef0: None,
        };
        let mut all_results = HashMap::new();

        if self.config.verbose {
            println!(
                "Grid search over {} parameter combinations",
                gamma_values.len() * n_components_values.len()
            );
        }

        for &gamma in gamma_values {
            for &n_components in &n_components_values {
                let params = ParameterSet {
                    gamma,
                    n_components,
                    degree: None,
                    coef0: None,
                };

                let cv_result = self.cross_validate_rbf(x, y, &params)?;
                let mean_score = cv_result.mean_test_score;

                all_results.insert(params.clone(), cv_result);

                if mean_score > best_score {
                    best_score = mean_score;
                    best_params = params;
                }

                if self.config.verbose {
                    println!(
                        "gamma={:.6}, n_components={}: score={:.6} ± {:.6}",
                        gamma,
                        n_components,
                        mean_score,
                        all_results
                            .get(&ParameterSet {
                                gamma,
                                n_components,
                                degree: None,
                                coef0: None
                            })
                            .expect("operation should succeed")
                            .std_test_score
                    );
                }
            }
        }

        Ok((best_params, best_score, all_results))
    }

    fn create_splitter(&self) -> Result<Box<dyn CVSplitter + Send + Sync>> {
        match &self.config.cv_strategy {
            CVStrategy::KFold { n_folds, shuffle } => Ok(Box::new(KFoldSplitter::new(
                *n_folds,
                *shuffle,
                self.config.random_seed,
            ))),
            CVStrategy::TimeSeriesSplit {
                n_splits,
                max_train_size,
            } => Ok(Box::new(TimeSeriesSplitter::new(
                *n_splits,
                *max_train_size,
            ))),
            CVStrategy::MonteCarlo {
                n_splits,
                test_size,
            } => Ok(Box::new(MonteCarloCVSplitter::new(
                *n_splits,
                *test_size,
                self.config.random_seed,
            ))),
            _ => {
                // Fallback to K-fold for unsupported strategies
                Ok(Box::new(KFoldSplitter::new(
                    5,
                    true,
                    self.config.random_seed,
                )))
            }
        }
    }

    fn extract_samples(&self, x: &Array2<f64>, indices: &[usize]) -> Array2<f64> {
        let n_features = x.ncols();
        let mut result = Array2::zeros((indices.len(), n_features));

        for (i, &idx) in indices.iter().enumerate() {
            result.row_mut(i).assign(&x.row(idx));
        }

        result
    }

    fn extract_targets(&self, y: &Array1<f64>, indices: &[usize]) -> Array1<f64> {
        let mut result = Array1::zeros(indices.len());

        for (i, &idx) in indices.iter().enumerate() {
            result[i] = y[idx];
        }

        result
    }

    fn compute_score(
        &self,
        x: &Array2<f64>,
        x_transformed: &Array2<f64>,
        y: Option<&Array1<f64>>,
        gamma: f64,
    ) -> Result<f64> {
        match &self.config.scoring_metric {
            ScoringMetric::KernelAlignment => {
                self.compute_kernel_alignment(x, x_transformed, gamma)
            }
            ScoringMetric::MeanSquaredError => {
                if let Some(y_data) = y {
                    self.compute_mse(x_transformed, y_data)
                } else {
                    Err(SklearsError::InvalidInput(
                        "Target values required for MSE".to_string(),
                    ))
                }
            }
            ScoringMetric::MeanAbsoluteError => {
                if let Some(y_data) = y {
                    self.compute_mae(x_transformed, y_data)
                } else {
                    Err(SklearsError::InvalidInput(
                        "Target values required for MAE".to_string(),
                    ))
                }
            }
            ScoringMetric::R2Score => {
                if let Some(y_data) = y {
                    self.compute_r2_score(x_transformed, y_data)
                } else {
                    Err(SklearsError::InvalidInput(
                        "Target values required for R²".to_string(),
                    ))
                }
            }
            _ => {
                // Fallback to kernel alignment
                self.compute_kernel_alignment(x, x_transformed, gamma)
            }
        }
    }

    fn compute_kernel_alignment(
        &self,
        x: &Array2<f64>,
        x_transformed: &Array2<f64>,
        gamma: f64,
    ) -> Result<f64> {
        let n_samples = x.nrows().min(50); // Limit for efficiency
        let x_subset = x.slice(s![..n_samples, ..]);

        // Compute exact kernel matrix
        let mut k_exact = Array2::zeros((n_samples, n_samples));
        for i in 0..n_samples {
            for j in 0..n_samples {
                let diff = &x_subset.row(i) - &x_subset.row(j);
                let squared_norm = diff.dot(&diff);
                k_exact[[i, j]] = (-gamma * squared_norm).exp();
            }
        }

        // Compute approximate kernel matrix
        let x_transformed_subset = x_transformed.slice(s![..n_samples, ..]);
        let k_approx = x_transformed_subset.dot(&x_transformed_subset.t());

        // Compute alignment
        let k_exact_frobenius = k_exact.iter().map(|&x| x * x).sum::<f64>().sqrt();
        let k_approx_frobenius = k_approx.iter().map(|&x| x * x).sum::<f64>().sqrt();
        let k_product = (&k_exact * &k_approx).sum();

        let alignment = k_product / (k_exact_frobenius * k_approx_frobenius);
        Ok(alignment)
    }

    fn compute_mse(&self, _x_transformed: &Array2<f64>, y: &Array1<f64>) -> Result<f64> {
        // Simple linear regression MSE
        // In practice, you'd want to use a proper regressor
        let y_mean = y.mean().unwrap_or(0.0);
        let mse = y.iter().map(|&yi| (yi - y_mean).powi(2)).sum::<f64>() / y.len() as f64;
        Ok(-mse) // Negative because we want to maximize score
    }

    fn compute_mae(&self, _x_transformed: &Array2<f64>, y: &Array1<f64>) -> Result<f64> {
        // Simple linear regression MAE
        let y_mean = y.mean().unwrap_or(0.0);
        let mae = y.iter().map(|&yi| (yi - y_mean).abs()).sum::<f64>() / y.len() as f64;
        Ok(-mae) // Negative because we want to maximize score
    }

    fn compute_r2_score(&self, _x_transformed: &Array2<f64>, y: &Array1<f64>) -> Result<f64> {
        // Simple R² score computation
        let y_mean = y.mean().unwrap_or(0.0);
        let ss_tot = y.iter().map(|&yi| (yi - y_mean).powi(2)).sum::<f64>();
        let ss_res = y.iter().map(|&yi| (yi - y_mean).powi(2)).sum::<f64>(); // Simplified

        let r2 = 1.0 - (ss_res / ss_tot);
        Ok(r2)
    }

    fn aggregate_results(
        &self,
        fold_results: Vec<(f64, Option<f64>, f64, f64)>,
    ) -> Result<CrossValidationResult> {
        let test_scores: Vec<f64> = fold_results.iter().map(|(score, _, _, _)| *score).collect();
        let train_scores: Option<Vec<f64>> = if self.config.return_train_score {
            Some(
                fold_results
                    .iter()
                    .filter_map(|(_, train_score, _, _)| *train_score)
                    .collect(),
            )
        } else {
            None
        };
        let fit_times: Vec<f64> = fold_results
            .iter()
            .map(|(_, _, fit_time, _)| *fit_time)
            .collect();
        let score_times: Vec<f64> = fold_results
            .iter()
            .map(|(_, _, _, score_time)| *score_time)
            .collect();

        let mean_test_score = test_scores.iter().sum::<f64>() / test_scores.len() as f64;
        let variance_test = test_scores
            .iter()
            .map(|&score| (score - mean_test_score).powi(2))
            .sum::<f64>()
            / test_scores.len() as f64;
        let std_test_score = variance_test.sqrt();

        let (mean_train_score, std_train_score) = if let Some(ref train_scores) = train_scores {
            let mean = train_scores.iter().sum::<f64>() / train_scores.len() as f64;
            let variance = train_scores
                .iter()
                .map(|&score| (score - mean).powi(2))
                .sum::<f64>()
                / train_scores.len() as f64;
            (Some(mean), Some(variance.sqrt()))
        } else {
            (None, None)
        };

        Ok(CrossValidationResult {
            test_scores,
            train_scores,
            mean_test_score,
            std_test_score,
            mean_train_score,
            std_train_score,
            fit_times,
            score_times,
        })
    }
}

#[allow(non_snake_case)]
#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_abs_diff_eq;

    #[test]
    fn test_kfold_splitter() {
        let x = Array2::from_shape_vec((20, 3), (0..60).map(|i| i as f64).collect())
            .expect("operation should succeed");

        let splitter = KFoldSplitter::new(4, false, Some(42));
        let splits = splitter.split(&x, None);

        assert_eq!(splits.len(), 4);

        // Check that all indices are covered exactly once
        let mut all_test_indices: Vec<usize> = Vec::new();
        for (_, test_indices) in &splits {
            all_test_indices.extend(test_indices);
        }
        all_test_indices.sort();

        let expected_indices: Vec<usize> = (0..20).collect();
        assert_eq!(all_test_indices, expected_indices);

        // Check fold sizes are approximately equal
        for (_, test_indices) in &splits {
            assert!(test_indices.len() >= 4);
            assert!(test_indices.len() <= 6);
        }
    }

    #[test]
    fn test_time_series_splitter() {
        let x = Array2::from_shape_vec((30, 2), (0..60).map(|i| i as f64).collect())
            .expect("operation should succeed");

        let splitter = TimeSeriesSplitter::new(3, Some(15));
        let splits = splitter.split(&x, None);

        assert_eq!(splits.len(), 3);

        // Check that training sets are chronologically before test sets
        for (train_indices, test_indices) in &splits {
            if !train_indices.is_empty() && !test_indices.is_empty() {
                let max_train = train_indices
                    .iter()
                    .max()
                    .expect("operation should succeed");
                let min_test = test_indices.iter().min().expect("operation should succeed");
                assert!(max_train < min_test);
            }
        }
    }

    #[test]
    fn test_monte_carlo_splitter() {
        let x = Array2::from_shape_vec((50, 4), (0..200).map(|i| i as f64).collect())
            .expect("operation should succeed");

        let splitter = MonteCarloCVSplitter::new(5, 0.3, Some(123));
        let splits = splitter.split(&x, None);

        assert_eq!(splits.len(), 5);

        // Check test size is approximately correct
        for (train_indices, test_indices) in &splits {
            let total_size = train_indices.len() + test_indices.len();
            assert_eq!(total_size, 50);
            assert!(test_indices.len() >= 14); // 30% of 50 = 15, allow some variance
            assert!(test_indices.len() <= 16);
        }
    }

    #[test]
    fn test_cross_validator_rbf() {
        let x = Array2::from_shape_vec((40, 5), (0..200).map(|i| i as f64 * 0.01).collect())
            .expect("operation should succeed");

        let config = CrossValidationConfig {
            cv_strategy: CVStrategy::KFold {
                n_folds: 3,
                shuffle: true,
            },
            scoring_metric: ScoringMetric::KernelAlignment,
            return_train_score: true,
            random_seed: Some(42),
            ..Default::default()
        };

        let cv = CrossValidator::new(config);
        let params = ParameterSet {
            gamma: 0.5,
            n_components: 20,
            degree: None,
            coef0: None,
        };

        let result = cv
            .cross_validate_rbf(&x, None, &params)
            .expect("operation should succeed");

        assert_eq!(result.test_scores.len(), 3);
        assert!(result.train_scores.is_some());
        assert_eq!(
            result
                .train_scores
                .as_ref()
                .expect("operation should succeed")
                .len(),
            3
        );
        assert!(result.mean_test_score > 0.0);
        assert!(result.std_test_score >= 0.0);
        assert!(result.mean_train_score.is_some());
        assert!(result.std_train_score.is_some());
        assert_eq!(result.fit_times.len(), 3);
        assert_eq!(result.score_times.len(), 3);
    }

    #[test]
    fn test_cross_validator_nystroem() {
        let x = Array2::from_shape_vec((30, 4), (0..120).map(|i| i as f64 * 0.02).collect())
            .expect("operation should succeed");

        let config = CrossValidationConfig {
            cv_strategy: CVStrategy::KFold {
                n_folds: 4,
                shuffle: false,
            },
            scoring_metric: ScoringMetric::KernelAlignment,
            ..Default::default()
        };

        let cv = CrossValidator::new(config);
        let params = ParameterSet {
            gamma: 1.0,
            n_components: 15,
            degree: None,
            coef0: None,
        };

        let result = cv
            .cross_validate_nystroem(&x, None, &params)
            .expect("operation should succeed");

        assert_eq!(result.test_scores.len(), 4);
        assert!(result.mean_test_score > 0.0);
        assert!(result.std_test_score >= 0.0);
    }

    #[test]
    fn test_grid_search_cv() {
        let x = Array2::from_shape_vec((25, 3), (0..75).map(|i| i as f64 * 0.05).collect())
            .expect("operation should succeed");

        let config = CrossValidationConfig {
            cv_strategy: CVStrategy::KFold {
                n_folds: 3,
                shuffle: true,
            },
            random_seed: Some(789),
            verbose: false,
            ..Default::default()
        };

        let cv = CrossValidator::new(config);

        let mut param_grid = HashMap::new();
        param_grid.insert("gamma".to_string(), vec![0.1, 1.0]);
        param_grid.insert("n_components".to_string(), vec![10.0, 20.0]);

        let (best_params, best_score, all_results) = cv
            .grid_search_cv(&x, None, &param_grid)
            .expect("operation should succeed");

        assert!(best_score > 0.0);
        assert!(best_params.gamma == 0.1 || best_params.gamma == 1.0);
        assert!(best_params.n_components == 10 || best_params.n_components == 20);
        assert_eq!(all_results.len(), 4); // 2x2 grid

        // Verify that best_score is actually the maximum
        let max_score = all_results
            .values()
            .map(|result| result.mean_test_score)
            .fold(f64::NEG_INFINITY, f64::max);
        assert_abs_diff_eq!(best_score, max_score, epsilon = 1e-10);
    }

    #[test]
    fn test_cross_validation_with_targets() {
        let x = Array2::from_shape_vec((20, 3), (0..60).map(|i| i as f64 * 0.1).collect())
            .expect("operation should succeed");
        let y = Array1::from_shape_fn(20, |i| (i as f64 * 0.1).sin());

        let config = CrossValidationConfig {
            cv_strategy: CVStrategy::KFold {
                n_folds: 4,
                shuffle: true,
            },
            scoring_metric: ScoringMetric::MeanSquaredError,
            random_seed: Some(456),
            ..Default::default()
        };

        let cv = CrossValidator::new(config);
        let params = ParameterSet {
            gamma: 0.8,
            n_components: 15,
            degree: None,
            coef0: None,
        };

        let result = cv
            .cross_validate_rbf(&x, Some(&y), &params)
            .expect("operation should succeed");

        assert_eq!(result.test_scores.len(), 4);
        // MSE scores are negative (we negate them to convert to maximization)
        assert!(result.mean_test_score <= 0.0);
    }

    #[test]
    fn test_cv_splitter_consistency() {
        let x = Array2::from_shape_vec((15, 2), (0..30).map(|i| i as f64).collect())
            .expect("operation should succeed");

        // Test that the same splitter with same seed produces same results
        let splitter1 = KFoldSplitter::new(3, true, Some(42));
        let splitter2 = KFoldSplitter::new(3, true, Some(42));

        let splits1 = splitter1.split(&x, None);
        let splits2 = splitter2.split(&x, None);

        assert_eq!(splits1.len(), splits2.len());
        for (split1, split2) in splits1.iter().zip(splits2.iter()) {
            assert_eq!(split1.0, split2.0); // train indices
            assert_eq!(split1.1, split2.1); // test indices
        }
    }

    #[test]
    fn test_cross_validation_result_aggregation() {
        let config = CrossValidationConfig {
            return_train_score: true,
            ..Default::default()
        };
        let cv = CrossValidator::new(config);

        let fold_results = vec![
            (0.8, Some(0.85), 0.1, 0.05),
            (0.75, Some(0.8), 0.12, 0.04),
            (0.82, Some(0.88), 0.11, 0.06),
        ];

        let result = cv
            .aggregate_results(fold_results)
            .expect("operation should succeed");

        assert_abs_diff_eq!(result.mean_test_score, 0.79, epsilon = 1e-10);
        assert!(result.std_test_score > 0.0);
        assert!(result.mean_train_score.is_some());
        assert_abs_diff_eq!(
            result.mean_train_score.expect("operation should succeed"),
            0.8433333333333334,
            epsilon = 1e-10
        );
    }
}
