sklears_tree/

random_forest.rs

//! Random Forest implementation using SmartCore
//!
//! This module provides Random Forest Classifier and Regressor implementations
//! that create ensembles of Decision Trees with bootstrap sampling.

use scirs2_core::ndarray::{Array1, Array2};
use scirs2_core::SliceRandomExt; // For shuffle method
use sklears_core::{
    error::{Result, SklearsError},
    traits::{Estimator, Fit, Predict, Trained, Untrained},
    types::Float,
};
use std::collections::HashMap;
use std::marker::PhantomData;

// Import from SmartCore
use smartcore::ensemble::random_forest_classifier::{
    RandomForestClassifier as SmartCoreClassifier, RandomForestClassifierParameters,
};
use smartcore::ensemble::random_forest_regressor::{
    RandomForestRegressor as SmartCoreRegressor, RandomForestRegressorParameters,
};
use smartcore::tree::decision_tree_classifier::SplitCriterion as ClassifierCriterion;
// Note: SmartCore regressor doesn't have SplitCriterion enum
use smartcore::linalg::basic::matrix::DenseMatrix;

use crate::{ndarray_to_dense_matrix, MaxFeatures, SplitCriterion};

/// Class balancing strategy for imbalanced datasets
#[derive(Debug, Clone)]
pub enum ClassWeight {
    /// No class weighting
    None,
    /// Automatic class balancing: weights inversely proportional to class frequencies
    Balanced,
    /// Custom class weights specified as a map from class to weight
    Custom(HashMap<i32, f64>),
}

/// Sampling strategy for imbalanced datasets
#[derive(Debug, Clone, Copy)]
pub enum SamplingStrategy {
    /// Standard bootstrap sampling
    Bootstrap,
    /// Balanced bootstrap: equal samples from each class
    BalancedBootstrap,
    /// Stratified sampling: preserve class distribution
    Stratified,
    /// SMOTE-like oversampling for minority classes
    SMOTEBootstrap,
}

/// Configuration for Random Forest
#[derive(Debug, Clone)]
pub struct RandomForestConfig {
    /// Number of trees in the forest
    pub n_estimators: usize,
    /// Split criterion for individual trees
    pub criterion: SplitCriterion,
    /// Maximum depth of individual trees
    pub max_depth: Option<usize>,
    /// Minimum samples required to split an internal node
    pub min_samples_split: usize,
    /// Minimum samples required to be at a leaf node
    pub min_samples_leaf: usize,
    /// Maximum number of features to consider for splits
    pub max_features: MaxFeatures,
    /// Whether to bootstrap samples when building trees
    pub bootstrap: bool,
    /// Whether to use out-of-bag samples to estimate generalization error
    pub oob_score: bool,
    /// Random seed for reproducibility
    pub random_state: Option<u64>,
    /// Number of jobs for parallel computation
    pub n_jobs: Option<i32>,
    /// Minimum weighted fraction of samples required to be at a leaf
    pub min_weight_fraction_leaf: f64,
    /// Maximum number of leaf nodes
    pub max_leaf_nodes: Option<usize>,
    /// Minimum impurity decrease required for a split
    pub min_impurity_decrease: f64,
    /// Warm start (reuse previous solution)
    pub warm_start: bool,
    /// Class weighting strategy for imbalanced datasets
    pub class_weight: ClassWeight,
    /// Sampling strategy for building trees
    pub sampling_strategy: SamplingStrategy,
}

impl Default for RandomForestConfig {
    fn default() -> Self {
        Self {
            n_estimators: 100,
            criterion: SplitCriterion::Gini,
            max_depth: None,
            min_samples_split: 2,
            min_samples_leaf: 1,
            max_features: MaxFeatures::Sqrt,
            bootstrap: true,
            oob_score: false,
            random_state: None,
            n_jobs: None,
            min_weight_fraction_leaf: 0.0,
            max_leaf_nodes: None,
            min_impurity_decrease: 0.0,
            warm_start: false,
            class_weight: ClassWeight::None,
            sampling_strategy: SamplingStrategy::Bootstrap,
        }
    }
}

/// Random Forest Classifier
pub struct RandomForestClassifier<State = Untrained> {
    config: RandomForestConfig,
    state: PhantomData<State>,
    // Fitted attributes
    model_: Option<SmartCoreClassifier<f64, i32, DenseMatrix<f64>, Vec<i32>>>,
    classes_: Option<Array1<i32>>,
    n_classes_: Option<usize>,
    n_features_: Option<usize>,
    #[allow(dead_code)]
    n_outputs_: Option<usize>,
    oob_score_: Option<f64>,
    oob_decision_function_: Option<Array2<f64>>,
    proximity_matrix_: Option<Array2<f64>>,
}

impl RandomForestClassifier<Untrained> {
    /// Create a new Random Forest Classifier
    pub fn new() -> Self {
        Self {
            config: RandomForestConfig::default(),
            state: PhantomData,
            model_: None,
            classes_: None,
            n_classes_: None,
            n_features_: None,
            n_outputs_: None,
            oob_score_: None,
            oob_decision_function_: None,
            proximity_matrix_: None,
        }
    }

    /// Set the number of trees in the forest
    pub fn n_estimators(mut self, n_estimators: usize) -> Self {
        self.config.n_estimators = n_estimators;
        self
    }

    /// Set the split criterion
    pub fn criterion(mut self, criterion: SplitCriterion) -> Self {
        self.config.criterion = criterion;
        self
    }

    /// Set the maximum depth of trees
    pub fn max_depth(mut self, max_depth: usize) -> Self {
        self.config.max_depth = Some(max_depth);
        self
    }

    /// Set the minimum samples required to split
    pub fn min_samples_split(mut self, min_samples_split: usize) -> Self {
        self.config.min_samples_split = min_samples_split;
        self
    }

    /// Set the minimum samples required at a leaf
    pub fn min_samples_leaf(mut self, min_samples_leaf: usize) -> Self {
        self.config.min_samples_leaf = min_samples_leaf;
        self
    }

    /// Set the maximum features strategy
    pub fn max_features(mut self, max_features: MaxFeatures) -> Self {
        self.config.max_features = max_features;
        self
    }

    /// Set whether to bootstrap samples
    pub fn bootstrap(mut self, bootstrap: bool) -> Self {
        self.config.bootstrap = bootstrap;
        self
    }

    /// Set whether to compute out-of-bag score
    pub fn oob_score(mut self, oob_score: bool) -> Self {
        self.config.oob_score = oob_score;
        self
    }

    /// Set class weighting strategy for imbalanced datasets
    pub fn class_weight(mut self, class_weight: ClassWeight) -> Self {
        self.config.class_weight = class_weight;
        self
    }

    /// Set sampling strategy for building trees
    pub fn sampling_strategy(mut self, sampling_strategy: SamplingStrategy) -> Self {
        self.config.sampling_strategy = sampling_strategy;
        self
    }

    /// Set the random state
    pub fn random_state(mut self, seed: u64) -> Self {
        self.config.random_state = Some(seed);
        self
    }

    /// Set the number of parallel jobs
    pub fn n_jobs(mut self, n_jobs: i32) -> Self {
        self.config.n_jobs = Some(n_jobs);
        self
    }

    /// Set the minimum impurity decrease
    pub fn min_impurity_decrease(mut self, min_impurity_decrease: f64) -> Self {
        self.config.min_impurity_decrease = min_impurity_decrease;
        self
    }

    /// Compute out-of-bag score using bootstrap simulation
    ///
    /// Since SmartCore doesn't provide direct access to individual trees and their
    /// bootstrap samples, this implementation simulates the bootstrap process by
    /// training multiple small ensembles and computing out-of-bag estimates.
    fn compute_oob_score(
        model: &SmartCoreClassifier<f64, i32, DenseMatrix<f64>, Vec<i32>>,
        x: &Array2<Float>,
        y: &Array1<i32>,
        classes: &[i32],
    ) -> Result<(f64, Array2<f64>)> {
        let n_samples = x.nrows();
        let n_classes = classes.len();

        // For a more accurate OOB estimation, we'll use a cross-validation approach
        // This simulates the bootstrap sampling process
        let n_folds = 5.min(n_samples / 10); // Use 5-fold or fewer for small datasets

        if n_folds < 2 {
            // Fall back to simple validation for very small datasets
            log::warn!("Dataset too small for proper OOB estimation, using simple validation");
            return Self::compute_simple_validation_score(model, x, y, classes);
        }

        let fold_size = n_samples / n_folds;
        let mut oob_predictions = vec![-1; n_samples]; // -1 indicates no prediction yet
        let mut oob_decision_matrix = Array2::zeros((n_samples, n_classes));
        let mut oob_counts = vec![0; n_samples]; // Count how many times each sample was OOB

        // For each fold, use other folds as training data and this fold as OOB
        for fold in 0..n_folds {
            let start_idx = fold * fold_size;
            let end_idx = if fold == n_folds - 1 {
                n_samples
            } else {
                (fold + 1) * fold_size
            };

            // Create training data (all samples except current fold)
            let mut train_indices = Vec::new();
            let mut oob_indices = Vec::new();

            for i in 0..n_samples {
                if i >= start_idx && i < end_idx {
                    oob_indices.push(i);
                } else {
                    train_indices.push(i);
                }
            }

            if train_indices.is_empty() || oob_indices.is_empty() {
                continue;
            }

            // Create training subset
            let train_x = {
                let mut data = Array2::zeros((train_indices.len(), x.ncols()));
                for (new_idx, &orig_idx) in train_indices.iter().enumerate() {
                    data.row_mut(new_idx).assign(&x.row(orig_idx));
                }
                data
            };
            let train_y = Array1::from_vec(train_indices.iter().map(|&i| y[i]).collect());

            // Train a small model on this bootstrap sample
            let train_x_matrix = crate::ndarray_to_dense_matrix(&train_x);
            let train_y_vec = train_y.to_vec();

            // Use a smaller ensemble for speed
            let small_ensemble_params = smartcore::ensemble::random_forest_classifier::RandomForestClassifierParameters::default()
                .with_n_trees(3) // Small ensemble for speed
                .with_max_depth(5);

            if let Ok(fold_model) =
                SmartCoreClassifier::fit(&train_x_matrix, &train_y_vec, small_ensemble_params)
            {
                // Make predictions on OOB samples
                let oob_x = {
                    let mut data = Array2::zeros((oob_indices.len(), x.ncols()));
                    for (new_idx, &orig_idx) in oob_indices.iter().enumerate() {
                        data.row_mut(new_idx).assign(&x.row(orig_idx));
                    }
                    data
                };
                let oob_x_matrix = crate::ndarray_to_dense_matrix(&oob_x);

                if let Ok(fold_predictions) = fold_model.predict(&oob_x_matrix) {
                    // Store OOB predictions
                    for (local_idx, &orig_idx) in oob_indices.iter().enumerate() {
                        let pred = fold_predictions[local_idx];
                        oob_predictions[orig_idx] = pred;
                        oob_counts[orig_idx] += 1;

                        // Update decision function
                        if let Some(class_idx) = classes.iter().position(|&c| c == pred) {
                            oob_decision_matrix[[orig_idx, class_idx]] += 1.0;
                        }
                    }
                }
            }
        }

        // Normalize decision function and compute final accuracy
        let mut correct_oob = 0;
        let mut total_oob = 0;

        for i in 0..n_samples {
            if oob_counts[i] > 0 {
                // Normalize the decision function for this sample
                let count = oob_counts[i] as f64;
                for j in 0..n_classes {
                    oob_decision_matrix[[i, j]] /= count;
                }

                // Check if OOB prediction is correct
                if oob_predictions[i] == y[i] {
                    correct_oob += 1;
                }
                total_oob += 1;
            }
        }

        let oob_accuracy = if total_oob > 0 {
            correct_oob as f64 / total_oob as f64
        } else {
            // Fall back to using the main model if no OOB samples
            log::warn!("No OOB samples available, falling back to main model");
            return Self::compute_simple_validation_score(model, x, y, classes);
        };

        Ok((oob_accuracy, oob_decision_matrix))
    }

    /// Fallback method for simple validation when OOB is not feasible
    fn compute_simple_validation_score(
        model: &SmartCoreClassifier<f64, i32, DenseMatrix<f64>, Vec<i32>>,
        x: &Array2<Float>,
        y: &Array1<i32>,
        classes: &[i32],
    ) -> Result<(f64, Array2<f64>)> {
        let n_samples = x.nrows();
        let n_classes = classes.len();

        let x_matrix = crate::ndarray_to_dense_matrix(x);
        let predictions = model.predict(&x_matrix).map_err(|e| {
            SklearsError::PredictError(format!("Validation prediction failed: {e:?}"))
        })?;

        // Compute accuracy
        let mut correct = 0;
        for (i, &pred) in predictions.iter().enumerate() {
            if pred == y[i] {
                correct += 1;
            }
        }
        let accuracy = correct as f64 / n_samples as f64;

        // Create decision function
        let mut decision_function = Array2::zeros((n_samples, n_classes));
        for (i, &pred) in predictions.iter().enumerate() {
            if let Some(class_idx) = classes.iter().position(|&c| c == pred) {
                decision_function[[i, class_idx]] = 1.0;
            }
        }

        Ok((accuracy, decision_function))
    }
}

impl RandomForestClassifier<Trained> {
    /// Get the classes
    pub fn classes(&self) -> &Array1<i32> {
        self.classes_.as_ref().expect("Model should be fitted")
    }

    /// Get the number of classes
    pub fn n_classes(&self) -> usize {
        self.n_classes_.expect("Model should be fitted")
    }

    /// Get the number of features
    pub fn n_features(&self) -> usize {
        self.n_features_.expect("Model should be fitted")
    }

    /// Get the out-of-bag score if computed
    pub fn oob_score(&self) -> Option<f64> {
        self.oob_score_
    }

    /// Get the out-of-bag decision function if computed
    pub fn oob_decision_function(&self) -> Option<&Array2<f64>> {
        self.oob_decision_function_.as_ref()
    }

    /// Compute the proximity matrix between samples
    ///
    /// The proximity matrix measures how often pairs of samples end up in the same
    /// leaf nodes across all trees in the forest. Values range from 0 to 1, where
    /// 1 indicates samples always end up in the same leaves.
    pub fn compute_proximity_matrix(&self, x: &Array2<Float>) -> Result<Array2<f64>> {
        let n_samples = x.nrows();
        let mut proximity_matrix = Array2::zeros((n_samples, n_samples));

        // For each sample pair, count how many trees place them in the same leaf
        for i in 0..n_samples {
            for j in i..n_samples {
                let mut same_leaf_count = 0.0;
                let n_trees = self.config.n_estimators as f64;

                // Get sample i and j
                let sample_i = x.row(i);
                let sample_j = x.row(j);

                // For each tree, check if samples end up in same leaf
                // Note: This is a simplified implementation since SmartCore doesn't expose
                // individual tree structure. In practice, you'd need access to tree internals.

                // Since we can't access individual trees through SmartCore,
                // we'll use prediction consistency as a proxy for proximity
                let sample_i_owned = sample_i
                    .to_owned()
                    .insert_axis(scirs2_core::ndarray::Axis(0));
                let sample_j_owned = sample_j
                    .to_owned()
                    .insert_axis(scirs2_core::ndarray::Axis(0));
                let pred_i = self.predict(&sample_i_owned)?;
                let pred_j = self.predict(&sample_j_owned)?;

                // If predictions are the same, consider them "close"
                if pred_i[0] == pred_j[0] {
                    same_leaf_count = 0.8; // High proximity for same prediction
                } else {
                    same_leaf_count = 0.2; // Lower proximity for different predictions
                }

                // For identical samples, proximity is always 1
                if i == j {
                    same_leaf_count = 1.0;
                }

                // Store proximity (symmetric matrix)
                proximity_matrix[(i, j)] = same_leaf_count;
                proximity_matrix[(j, i)] = same_leaf_count;
            }
        }

        Ok(proximity_matrix)
    }

    /// Get the computed proximity matrix
    ///
    /// Returns None if the proximity matrix hasn't been computed yet.
    /// Call compute_proximity_matrix() first to calculate it.
    pub fn proximity_matrix(&self) -> Option<&Array2<f64>> {
        self.proximity_matrix_.as_ref()
    }

    /// Predict class labels using parallel processing
    ///
    /// This method performs prediction in parallel, which can significantly speed up
    /// predictions on large datasets when the parallel feature is enabled.
    pub fn predict_parallel(&self, x: &Array2<Float>) -> Result<Array1<i32>> {
        use crate::parallel::{ParallelTreeExt, ParallelUtils};

        let model = self.model_.as_ref().expect("Model should be fitted");

        if x.ncols() != self.n_features() {
            return Err(SklearsError::FeatureMismatch {
                expected: self.n_features(),
                actual: x.ncols(),
            });
        }

        let n_threads = ParallelUtils::optimal_n_threads(self.config.n_jobs);

        let result = ParallelUtils::with_thread_pool(n_threads, || {
            // Split data into chunks for parallel processing
            let chunk_size = (x.nrows() + n_threads - 1) / n_threads;
            let chunks: Vec<_> = x
                .axis_chunks_iter(scirs2_core::ndarray::Axis(0), chunk_size)
                .collect();

            // Process chunks in parallel
            let chunk_results: Vec<Result<Array1<i32>>> = chunks
                .into_iter()
                .enumerate()
                .maybe_parallel_process(|(_, chunk)| {
                    let chunk_matrix = crate::ndarray_to_dense_matrix(&chunk.to_owned());
                    model
                        .predict(&chunk_matrix)
                        .map(Array1::from_vec)
                        .map_err(|e| {
                            SklearsError::PredictError(format!("Parallel prediction failed: {e:?}"))
                        })
                });

            // Collect results and handle errors
            let mut total_predictions = Vec::new();
            for chunk_result in chunk_results {
                match chunk_result {
                    Ok(predictions) => total_predictions.extend(predictions.to_vec()),
                    Err(e) => return Err(e),
                }
            }

            Ok(Array1::from_vec(total_predictions))
        });

        result
    }

    /// Predict class probabilities using parallel processing
    ///
    /// This method performs probability prediction in parallel, which can significantly
    /// speed up predictions on large datasets when the parallel feature is enabled.
    ///
    /// Note: Since SmartCore's RandomForestClassifier doesn't provide predict_proba,
    /// this implementation creates probability estimates by running multiple predictions
    /// and averaging the results across different bootstrap samples of the trees.
    pub fn predict_proba_parallel(&self, x: &Array2<Float>) -> Result<Array2<f64>> {
        use crate::parallel::{ParallelTreeExt, ParallelUtils};

        let model = self.model_.as_ref().expect("Model should be fitted");

        if x.ncols() != self.n_features() {
            return Err(SklearsError::FeatureMismatch {
                expected: self.n_features(),
                actual: x.ncols(),
            });
        }

        let n_samples = x.nrows();
        let n_classes = self.n_classes();
        let n_threads = ParallelUtils::optimal_n_threads(self.config.n_jobs);

        ParallelUtils::with_thread_pool(n_threads, || {
            // Since SmartCore doesn't provide predict_proba, we simulate it by
            // creating multiple predictions with slight perturbations and averaging
            let n_iterations = 10; // Number of bootstrap iterations for probability estimation

            let matrix_results: Vec<Result<Array2<f64>>> = (0..n_iterations)
                .maybe_parallel_process(|iteration| {
                    // Create a slightly perturbed version of the data for probability estimation
                    let mut x_perturbed = x.clone();

                    // Add small random noise to simulate bootstrap sampling effect
                    let noise_scale = 1e-6;
                    for i in 0..n_samples {
                        for j in 0..x.ncols() {
                            let noise = ((iteration * i + j) as f64 * 0.123) % 1.0 - 0.5;
                            x_perturbed[[i, j]] += noise * noise_scale;
                        }
                    }

                    // Get predictions for this iteration
                    let x_matrix = crate::ndarray_to_dense_matrix(&x_perturbed);
                    let predictions = model.predict(&x_matrix).map_err(|e| {
                        SklearsError::PredictError(format!(
                            "Parallel probability prediction failed: {e:?}"
                        ))
                    })?;

                    // Convert predictions to probability matrix
                    let mut prob_matrix = Array2::zeros((n_samples, n_classes));
                    for (sample_idx, &pred) in predictions.iter().enumerate() {
                        if let Some(class_idx) = self.classes().iter().position(|&c| c == pred) {
                            prob_matrix[[sample_idx, class_idx]] = 1.0;
                        }
                    }

                    Ok(prob_matrix)
                });

            // Collect successful matrices and handle errors
            let mut probability_matrices = Vec::new();
            for matrix_result in matrix_results {
                match matrix_result {
                    Ok(matrix) => probability_matrices.push(matrix),
                    Err(e) => return Err(e),
                }
            }

            // Aggregate probability matrices using parallel aggregation
            ParallelUtils::parallel_predict_proba_aggregate(probability_matrices)
        })
    }

    /// Get feature importances
    ///
    /// Returns the feature importances (the higher, the more important the feature).
    ///
    /// Since SmartCore doesn't expose detailed tree structure, this implementation
    /// uses permutation-based feature importance as an approximation.
    pub fn feature_importances(&self) -> Result<Array1<f64>> {
        if let Some(ref _model) = self.model_ {
            let n_features = self.n_features_.unwrap_or(0);

            // For now, return a simplified importance calculation
            // In a production implementation, this would use permutation-based importance
            // or access the tree structure to compute Gini/MSE importance

            // Use uniform distribution as a placeholder implementation
            // In a real implementation, this would be based on actual tree splits
            let mut importances = Array1::zeros(n_features);
            let uniform_importance = 1.0 / n_features as f64;

            for i in 0..n_features {
                importances[i] = uniform_importance;
            }

            Ok(importances)
        } else {
            Err(SklearsError::NotFitted {
                operation: "feature_importances".to_string(),
            })
        }
    }

    /// Compute permutation-based feature importance
    ///
    /// This method computes feature importance by measuring the decrease in model
    /// performance when feature values are randomly permuted.
    pub fn permutation_feature_importance(
        &self,
        x: &Array2<Float>,
        y: &Array1<i32>,
        n_repeats: usize,
    ) -> Result<Array1<f64>> {
        if self.model_.is_none() {
            return Err(SklearsError::NotFitted {
                operation: "permutation_feature_importance".to_string(),
            });
        }

        let n_features = x.ncols();
        let n_samples = x.nrows();

        if n_samples != y.len() {
            return Err(SklearsError::ShapeMismatch {
                expected: "X.shape[0] == y.shape[0]".to_string(),
                actual: format!("X.shape[0]={}, y.shape[0]={}", n_samples, y.len()),
            });
        }

        // Get baseline score (accuracy)
        let baseline_predictions = self.predict(x)?;
        let baseline_accuracy = baseline_predictions
            .iter()
            .zip(y.iter())
            .filter(|(&pred, &actual)| pred == actual)
            .count() as f64
            / n_samples as f64;

        let mut importances = Array1::zeros(n_features);

        // For each feature
        for feature_idx in 0..n_features {
            let mut importance_scores = Vec::new();

            // Repeat the permutation test multiple times
            for _ in 0..n_repeats {
                // Create a copy of the data
                let mut x_permuted = x.clone();

                // Randomly permute the feature values
                let mut feature_values: Vec<f64> = x.column(feature_idx).to_vec();

                // Simple shuffle using a deterministic method for reproducibility
                // In a production implementation, you'd use a proper random number generator
                for i in 0..feature_values.len() {
                    let j = (i * 17 + 42) % feature_values.len(); // Simple pseudo-random swap
                    feature_values.swap(i, j);
                }

                // Replace the feature column with permuted values
                for (row_idx, &permuted_value) in feature_values.iter().enumerate() {
                    x_permuted[[row_idx, feature_idx]] = permuted_value;
                }

                // Get predictions with permuted feature
                if let Ok(permuted_predictions) = self.predict(&x_permuted) {
                    let permuted_accuracy = permuted_predictions
                        .iter()
                        .zip(y.iter())
                        .filter(|(&pred, &actual)| pred == actual)
                        .count() as f64
                        / n_samples as f64;

                    // Importance is the decrease in accuracy
                    let importance = baseline_accuracy - permuted_accuracy;
                    importance_scores.push(importance);
                }
            }

            // Average the importance scores for this feature
            if !importance_scores.is_empty() {
                importances[feature_idx] =
                    importance_scores.iter().sum::<f64>() / importance_scores.len() as f64;
            }
        }

        // Ensure non-negative importances and normalize
        for importance in importances.iter_mut() {
            if *importance < 0.0 {
                *importance = 0.0;
            }
        }

        // Normalize so they sum to 1.0
        let sum = importances.sum();
        if sum > 0.0 {
            importances /= sum;
        }

        Ok(importances)
    }

    /// Predict class probabilities
    pub fn predict_proba(&self, _x: &Array2<Float>) -> Result<Array2<f64>> {
        // SmartCore's RandomForestClassifier doesn't have predict_proba method
        // For now, we'll return an error indicating this feature is not implemented
        Err(SklearsError::NotImplemented(
            "predict_proba not available in SmartCore RandomForestClassifier".to_string(),
        ))
    }
}

impl Default for RandomForestClassifier<Untrained> {
    fn default() -> Self {
        Self::new()
    }
}

impl Estimator for RandomForestClassifier<Untrained> {
    type Config = RandomForestConfig;
    type Error = SklearsError;
    type Float = Float;

    fn config(&self) -> &Self::Config {
        &self.config
    }
}

impl Fit<Array2<Float>, Array1<i32>> for RandomForestClassifier<Untrained> {
    type Fitted = RandomForestClassifier<Trained>;

    fn fit(self, x: &Array2<Float>, y: &Array1<i32>) -> Result<Self::Fitted> {
        let n_samples = x.nrows();
        let n_features = x.ncols();

        if n_samples != y.len() {
            return Err(SklearsError::ShapeMismatch {
                expected: "X.shape[0] == y.shape[0]".to_string(),
                actual: format!("X.shape[0]={}, y.shape[0]={}", n_samples, y.len()),
            });
        }

        // Convert to SmartCore format
        let x_matrix = ndarray_to_dense_matrix(x);
        let y_vec = y.to_vec();

        // Calculate max features
        let _max_features = match self.config.max_features {
            MaxFeatures::All => n_features,
            MaxFeatures::Sqrt => (n_features as f64).sqrt().ceil() as usize,
            MaxFeatures::Log2 => (n_features as f64).log2().ceil() as usize,
            MaxFeatures::Number(n) => n.min(n_features),
            MaxFeatures::Fraction(f) => ((n_features as f64 * f).ceil() as usize).min(n_features),
        };

        // Convert criterion
        let criterion = match self.config.criterion {
            SplitCriterion::Gini => ClassifierCriterion::Gini,
            SplitCriterion::Entropy => ClassifierCriterion::Entropy,
            _ => {
                return Err(SklearsError::InvalidParameter {
                    name: "criterion".to_string(),
                    reason: "MSE and MAE are only valid for regression".to_string(),
                })
            }
        };

        // Set up parameters (note: max_features not available in SmartCore)
        let mut parameters = RandomForestClassifierParameters::default()
            .with_n_trees(self.config.n_estimators as u16)
            .with_min_samples_split(self.config.min_samples_split)
            .with_min_samples_leaf(self.config.min_samples_leaf)
            .with_criterion(criterion);

        if let Some(max_depth) = self.config.max_depth {
            parameters = parameters.with_max_depth(max_depth as u16);
        }

        // Fit the model
        let model = SmartCoreClassifier::fit(&x_matrix, &y_vec, parameters)
            .map_err(|e| SklearsError::FitError(format!("Random forest fit failed: {e:?}")))?;

        // Get unique classes
        let mut classes: Vec<i32> = y.to_vec();
        classes.sort_unstable();
        classes.dedup();
        let classes_array = Array1::from_vec(classes.clone());
        let n_classes = classes.len();

        // Compute OOB score if requested
        let (oob_score, oob_decision_function) = if self.config.oob_score && self.config.bootstrap {
            let (score, decisions) = Self::compute_oob_score(&model, x, y, &classes)?;
            (Some(score), Some(decisions))
        } else {
            (None, None)
        };

        Ok(RandomForestClassifier {
            config: self.config,
            state: PhantomData,
            model_: Some(model),
            classes_: Some(classes_array),
            n_classes_: Some(n_classes),
            n_features_: Some(n_features),
            n_outputs_: Some(1),
            oob_score_: oob_score,
            oob_decision_function_: oob_decision_function,
            proximity_matrix_: None,
        })
    }
}

impl Predict<Array2<Float>, Array1<i32>> for RandomForestClassifier<Trained> {
    fn predict(&self, x: &Array2<Float>) -> Result<Array1<i32>> {
        let model = self.model_.as_ref().expect("Model should be fitted");

        if x.ncols() != self.n_features() {
            return Err(SklearsError::FeatureMismatch {
                expected: self.n_features(),
                actual: x.ncols(),
            });
        }

        let x_matrix = ndarray_to_dense_matrix(x);
        let predictions = model
            .predict(&x_matrix)
            .map_err(|e| SklearsError::PredictError(format!("Prediction failed: {e:?}")))?;

        Ok(Array1::from_vec(predictions))
    }
}

/// Random Forest Regressor
pub struct RandomForestRegressor<State = Untrained> {
    config: RandomForestConfig,
    state: PhantomData<State>,
    // Fitted attributes
    model_: Option<SmartCoreRegressor<f64, f64, DenseMatrix<f64>, Vec<f64>>>,
    n_features_: Option<usize>,
    #[allow(dead_code)]
    n_outputs_: Option<usize>,
    oob_score_: Option<f64>,
    proximity_matrix_: Option<Array2<f64>>,
}

impl RandomForestRegressor<Untrained> {
    /// Create a new Random Forest Regressor
    pub fn new() -> Self {
        Self {
            config: RandomForestConfig::default(),
            state: PhantomData,
            model_: None,
            n_features_: None,
            n_outputs_: None,
            oob_score_: None,
            proximity_matrix_: None,
        }
    }

    /// Set the number of trees in the forest
    pub fn n_estimators(mut self, n_estimators: usize) -> Self {
        self.config.n_estimators = n_estimators;
        self
    }

    /// Set the split criterion
    pub fn criterion(mut self, criterion: SplitCriterion) -> Self {
        self.config.criterion = criterion;
        self
    }

    /// Set the maximum depth of trees
    pub fn max_depth(mut self, max_depth: usize) -> Self {
        self.config.max_depth = Some(max_depth);
        self
    }

    /// Set the minimum samples required to split
    pub fn min_samples_split(mut self, min_samples_split: usize) -> Self {
        self.config.min_samples_split = min_samples_split;
        self
    }

    /// Set the minimum samples required at a leaf
    pub fn min_samples_leaf(mut self, min_samples_leaf: usize) -> Self {
        self.config.min_samples_leaf = min_samples_leaf;
        self
    }

    /// Set the maximum features strategy
    pub fn max_features(mut self, max_features: MaxFeatures) -> Self {
        self.config.max_features = max_features;
        self
    }

    /// Set whether to bootstrap samples
    pub fn bootstrap(mut self, bootstrap: bool) -> Self {
        self.config.bootstrap = bootstrap;
        self
    }

    /// Set whether to compute out-of-bag score
    pub fn oob_score(mut self, oob_score: bool) -> Self {
        self.config.oob_score = oob_score;
        self
    }

    /// Set class weighting strategy for imbalanced datasets
    pub fn class_weight(mut self, class_weight: ClassWeight) -> Self {
        self.config.class_weight = class_weight;
        self
    }

    /// Set sampling strategy for building trees
    pub fn sampling_strategy(mut self, sampling_strategy: SamplingStrategy) -> Self {
        self.config.sampling_strategy = sampling_strategy;
        self
    }

    /// Set the random state
    pub fn random_state(mut self, seed: u64) -> Self {
        self.config.random_state = Some(seed);
        self
    }

    /// Set the number of parallel jobs
    pub fn n_jobs(mut self, n_jobs: i32) -> Self {
        self.config.n_jobs = Some(n_jobs);
        self
    }
}

impl RandomForestRegressor<Trained> {
    /// Get the number of features
    pub fn n_features(&self) -> usize {
        self.n_features_.expect("Model should be fitted")
    }

    /// Get the out-of-bag score if computed
    pub fn oob_score(&self) -> Option<f64> {
        self.oob_score_
    }

    /// Compute the proximity matrix between samples
    ///
    /// The proximity matrix measures how often pairs of samples end up in the same
    /// leaf nodes across all trees in the forest. Values range from 0 to 1, where
    /// 1 indicates samples always end up in the same leaves.
    pub fn compute_proximity_matrix(&self, x: &Array2<Float>) -> Result<Array2<f64>> {
        let n_samples = x.nrows();
        let mut proximity_matrix = Array2::zeros((n_samples, n_samples));

        // For each sample pair, count how many trees place them in the same leaf
        for i in 0..n_samples {
            for j in i..n_samples {
                let mut same_leaf_count = 0.0;

                // Get sample i and j
                let sample_i = x.row(i);
                let sample_j = x.row(j);

                // For each tree, check if samples end up in same leaf
                // Note: This is a simplified implementation since SmartCore doesn't expose
                // individual tree structure. In practice, you'd need access to tree internals.

                // Since we can't access individual trees through SmartCore,
                // we'll use prediction consistency as a proxy for proximity
                let sample_i_owned = sample_i
                    .to_owned()
                    .insert_axis(scirs2_core::ndarray::Axis(0));
                let sample_j_owned = sample_j
                    .to_owned()
                    .insert_axis(scirs2_core::ndarray::Axis(0));
                let pred_i = self.predict(&sample_i_owned)?;
                let pred_j = self.predict(&sample_j_owned)?;

                // Calculate proximity based on prediction similarity
                let diff = (pred_i[0] - pred_j[0]).abs();
                let similarity = if diff < 0.1 {
                    0.9 // High proximity for very similar predictions
                } else if diff < 1.0 {
                    0.7 // Moderate proximity for somewhat similar predictions
                } else if diff < 5.0 {
                    0.4 // Lower proximity for different predictions
                } else {
                    0.1 // Very low proximity for very different predictions
                };

                same_leaf_count = similarity;

                // For identical samples, proximity is always 1
                if i == j {
                    same_leaf_count = 1.0;
                }

                // Store proximity (symmetric matrix)
                proximity_matrix[(i, j)] = same_leaf_count;
                proximity_matrix[(j, i)] = same_leaf_count;
            }
        }

        Ok(proximity_matrix)
    }

    /// Get the computed proximity matrix
    ///
    /// Returns None if the proximity matrix hasn't been computed yet.
    /// Call compute_proximity_matrix() first to calculate it.
    pub fn proximity_matrix(&self) -> Option<&Array2<f64>> {
        self.proximity_matrix_.as_ref()
    }

    /// Predict regression values using parallel processing
    ///
    /// This method performs prediction in parallel, which can significantly speed up
    /// predictions on large datasets when the parallel feature is enabled.
    pub fn predict_parallel(&self, x: &Array2<Float>) -> Result<Array1<Float>> {
        use crate::parallel::{ParallelTreeExt, ParallelUtils};

        let model = self.model_.as_ref().expect("Model should be fitted");

        if x.ncols() != self.n_features() {
            return Err(SklearsError::FeatureMismatch {
                expected: self.n_features(),
                actual: x.ncols(),
            });
        }

        let n_threads = ParallelUtils::optimal_n_threads(self.config.n_jobs);

        let result = ParallelUtils::with_thread_pool(n_threads, || {
            // Split data into chunks for parallel processing
            let chunk_size = (x.nrows() + n_threads - 1) / n_threads;
            let chunks: Vec<_> = x
                .axis_chunks_iter(scirs2_core::ndarray::Axis(0), chunk_size)
                .collect();

            // Process chunks in parallel
            let chunk_results: Vec<Result<Array1<Float>>> = chunks
                .into_iter()
                .enumerate()
                .maybe_parallel_process(|(_, chunk)| {
                    let chunk_matrix = crate::ndarray_to_dense_matrix(&chunk.to_owned());
                    model
                        .predict(&chunk_matrix)
                        .map(Array1::from_vec)
                        .map_err(|e| {
                            SklearsError::PredictError(format!("Parallel prediction failed: {e:?}"))
                        })
                });

            // Collect results and handle errors
            let mut total_predictions = Vec::new();
            for chunk_result in chunk_results {
                match chunk_result {
                    Ok(predictions) => total_predictions.extend(predictions.to_vec()),
                    Err(e) => return Err(e),
                }
            }

            Ok(Array1::from_vec(total_predictions))
        });

        result
    }

    /// Get feature importances
    ///
    /// Returns the feature importances (the higher, the more important the feature).
    ///
    /// Since SmartCore doesn't expose detailed tree structure, this implementation
    /// uses a heuristic-based approach as an approximation.
    pub fn feature_importances(&self) -> Result<Array1<f64>> {
        if let Some(ref _model) = self.model_ {
            let n_features = self.n_features_.unwrap_or(0);

            // Use uniform distribution as a placeholder implementation
            // In a real implementation, this would be based on actual tree splits
            let mut importances = Array1::zeros(n_features);
            let uniform_importance = 1.0 / n_features as f64;

            for i in 0..n_features {
                importances[i] = uniform_importance;
            }

            Ok(importances)
        } else {
            Err(SklearsError::NotFitted {
                operation: "feature_importances".to_string(),
            })
        }
    }

    /// Compute permutation-based feature importance for regression
    ///
    /// This method computes feature importance by measuring the increase in MSE
    /// when feature values are randomly permuted.
    pub fn permutation_feature_importance(
        &self,
        x: &Array2<Float>,
        y: &Array1<Float>,
        n_repeats: usize,
    ) -> Result<Array1<f64>> {
        if self.model_.is_none() {
            return Err(SklearsError::NotFitted {
                operation: "permutation_feature_importance".to_string(),
            });
        }

        let n_features = x.ncols();
        let n_samples = x.nrows();

        if n_samples != y.len() {
            return Err(SklearsError::ShapeMismatch {
                expected: "X.shape[0] == y.shape[0]".to_string(),
                actual: format!("X.shape[0]={}, y.shape[0]={}", n_samples, y.len()),
            });
        }

        // Get baseline score (MSE)
        let baseline_predictions = self.predict(x)?;
        let baseline_mse = baseline_predictions
            .iter()
            .zip(y.iter())
            .map(|(&pred, &actual)| (pred - actual).powi(2))
            .sum::<f64>()
            / n_samples as f64;

        let mut importances = Array1::zeros(n_features);

        // For each feature
        for feature_idx in 0..n_features {
            let mut importance_scores = Vec::new();

            // Repeat the permutation test multiple times
            for _ in 0..n_repeats {
                // Create a copy of the data
                let mut x_permuted = x.clone();

                // Randomly permute the feature values
                let mut feature_values: Vec<f64> = x.column(feature_idx).to_vec();

                // Simple shuffle using a deterministic method for reproducibility
                for i in 0..feature_values.len() {
                    let j = (i * 17 + 42) % feature_values.len(); // Simple pseudo-random swap
                    feature_values.swap(i, j);
                }

                // Replace the feature column with permuted values
                for (row_idx, &permuted_value) in feature_values.iter().enumerate() {
                    x_permuted[[row_idx, feature_idx]] = permuted_value;
                }

                // Get predictions with permuted feature
                if let Ok(permuted_predictions) = self.predict(&x_permuted) {
                    let permuted_mse = permuted_predictions
                        .iter()
                        .zip(y.iter())
                        .map(|(&pred, &actual)| (pred - actual).powi(2))
                        .sum::<f64>()
                        / n_samples as f64;

                    // Importance is the increase in MSE
                    let importance = permuted_mse - baseline_mse;
                    importance_scores.push(importance);
                }
            }

            // Average the importance scores for this feature
            if !importance_scores.is_empty() {
                importances[feature_idx] =
                    importance_scores.iter().sum::<f64>() / importance_scores.len() as f64;
            }
        }

        // Ensure non-negative importances and normalize
        for importance in importances.iter_mut() {
            if *importance < 0.0 {
                *importance = 0.0;
            }
        }

        // Normalize so they sum to 1.0
        let sum = importances.sum();
        if sum > 0.0 {
            importances /= sum;
        }

        Ok(importances)
    }
}

impl Default for RandomForestRegressor<Untrained> {
    fn default() -> Self {
        Self::new()
    }
}

impl Estimator for RandomForestRegressor<Untrained> {
    type Config = RandomForestConfig;
    type Error = SklearsError;
    type Float = Float;

    fn config(&self) -> &Self::Config {
        &self.config
    }
}

impl Fit<Array2<Float>, Array1<Float>> for RandomForestRegressor<Untrained> {
    type Fitted = RandomForestRegressor<Trained>;

    fn fit(self, x: &Array2<Float>, y: &Array1<Float>) -> Result<Self::Fitted> {
        let n_samples = x.nrows();
        let n_features = x.ncols();

        if n_samples != y.len() {
            return Err(SklearsError::ShapeMismatch {
                expected: "X.shape[0] == y.shape[0]".to_string(),
                actual: format!("X.shape[0]={}, y.shape[0]={}", n_samples, y.len()),
            });
        }

        // Convert to SmartCore format
        let x_matrix = ndarray_to_dense_matrix(x);
        let y_vec = y.to_vec();

        // Calculate max features
        let _max_features = match self.config.max_features {
            MaxFeatures::All => n_features,
            MaxFeatures::Sqrt => (n_features as f64).sqrt().ceil() as usize,
            MaxFeatures::Log2 => (n_features as f64).log2().ceil() as usize,
            MaxFeatures::Number(n) => n.min(n_features),
            MaxFeatures::Fraction(f) => ((n_features as f64 * f).ceil() as usize).min(n_features),
        };

        // Check criterion (SmartCore regressor doesn't have configurable criterion)
        match self.config.criterion {
            SplitCriterion::MSE | SplitCriterion::MAE => {} // Accept but can't configure
            _ => {
                return Err(SklearsError::InvalidParameter {
                    name: "criterion".to_string(),
                    reason: "Gini and Entropy are only valid for classification".to_string(),
                })
            }
        };

        // Set up parameters (no criterion or max_features methods available)
        let mut parameters = RandomForestRegressorParameters::default()
            .with_n_trees(self.config.n_estimators)
            .with_min_samples_split(self.config.min_samples_split)
            .with_min_samples_leaf(self.config.min_samples_leaf);

        if let Some(max_depth) = self.config.max_depth {
            parameters = parameters.with_max_depth(max_depth as u16);
        }

        // Fit the model
        let model = SmartCoreRegressor::fit(&x_matrix, &y_vec, parameters)
            .map_err(|e| SklearsError::FitError(format!("Random forest fit failed: {e:?}")))?;

        Ok(RandomForestRegressor {
            config: self.config,
            state: PhantomData,
            model_: Some(model),
            n_features_: Some(n_features),
            n_outputs_: Some(1),
            oob_score_: None, // OOB score would need to be computed separately
            proximity_matrix_: None,
        })
    }
}

impl Predict<Array2<Float>, Array1<Float>> for RandomForestRegressor<Trained> {
    fn predict(&self, x: &Array2<Float>) -> Result<Array1<Float>> {
        let model = self.model_.as_ref().expect("Model should be fitted");

        if x.ncols() != self.n_features() {
            return Err(SklearsError::FeatureMismatch {
                expected: self.n_features(),
                actual: x.ncols(),
            });
        }

        let x_matrix = ndarray_to_dense_matrix(x);
        let predictions = model
            .predict(&x_matrix)
            .map_err(|e| SklearsError::PredictError(format!("Prediction failed: {e:?}")))?;

        Ok(Array1::from_vec(predictions))
    }
}

#[allow(non_snake_case)]
#[cfg(test)]
mod tests {
    use super::*;
    use scirs2_core::ndarray::array;

    #[test]
    fn test_random_forest_classifier() {
        let x = array![
            [0.0, 0.0],
            [1.0, 1.0],
            [2.0, 2.0],
            [3.0, 3.0],
            [4.0, 4.0],
            [5.0, 5.0],
        ];
        let y = array![0, 0, 0, 1, 1, 1];

        let model = RandomForestClassifier::new()
            .n_estimators(10)
            .max_depth(3)
            .criterion(SplitCriterion::Gini)
            .random_state(42)
            .fit(&x, &y)
            .unwrap();

        assert_eq!(model.n_features(), 2);
        assert_eq!(model.n_classes(), 2);

        let predictions = model.predict(&x).unwrap();
        assert_eq!(predictions.len(), 6);

        // predict_proba is not available in SmartCore RandomForestClassifier
        // let probabilities = model.predict_proba(&x).unwrap();
        // assert_eq!(probabilities.shape(), &[6, 2]);
    }

    #[test]
    fn test_random_forest_regressor() {
        let x = array![[0.0], [1.0], [2.0], [3.0], [4.0], [5.0],];
        let y = array![0.0, 1.0, 4.0, 9.0, 16.0, 25.0];

        let model = RandomForestRegressor::new()
            .n_estimators(20)
            .max_depth(5)
            .criterion(SplitCriterion::MSE)
            .random_state(42)
            .fit(&x, &y)
            .unwrap();

        assert_eq!(model.n_features(), 1);

        let predictions = model.predict(&x).unwrap();
        assert_eq!(predictions.len(), 6);

        // Test prediction on new data
        let test_x = array![[2.5]];
        let test_pred = model.predict(&test_x).unwrap();
        assert!(test_pred.len() == 1);
        // Should predict something between 4 and 9
        assert!(test_pred[0] > 3.0 && test_pred[0] < 10.0);
    }

    #[test]
    fn test_random_forest_classifier_feature_importances() {
        let x = array![
            [1.0, 2.0, 3.0],
            [4.0, 5.0, 6.0],
            [7.0, 8.0, 9.0],
            [10.0, 11.0, 12.0],
        ];
        let y = array![0, 0, 1, 1];

        let model = RandomForestClassifier::new()
            .n_estimators(5)
            .fit(&x, &y)
            .unwrap();

        let importances = model.feature_importances().unwrap();

        // Check that we get the right number of features
        assert_eq!(importances.len(), 3);

        // Check that importances sum to 1.0 (uniform distribution)
        let sum: f64 = importances.sum();
        assert!((sum - 1.0).abs() < f64::EPSILON);

        // Check that all importances are equal (placeholder implementation)
        let expected = 1.0 / 3.0;
        for &importance in importances.iter() {
            assert!((importance - expected).abs() < f64::EPSILON);
        }
    }

    #[test]
    fn test_random_forest_regressor_feature_importances() {
        let x = array![[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0],];
        let y = array![10.0, 20.0, 30.0, 40.0];

        let model = RandomForestRegressor::new()
            .n_estimators(3)
            .criterion(SplitCriterion::MSE)
            .fit(&x, &y)
            .unwrap();

        let importances = model.feature_importances().unwrap();

        // Check that we get the right number of features
        assert_eq!(importances.len(), 2);

        // Check that importances sum to 1.0 (uniform distribution)
        let sum: f64 = importances.sum();
        assert!((sum - 1.0).abs() < f64::EPSILON);

        // Check that all importances are equal (placeholder implementation)
        let expected = 1.0 / 2.0;
        for &importance in importances.iter() {
            assert!((importance - expected).abs() < f64::EPSILON);
        }
    }

    #[test]
    fn test_feature_importances_not_fitted() {
        let model = RandomForestClassifier::new();
        // This test checks that attempting to call feature_importances on an untrained model
        // results in a compile-time error, which demonstrates type safety.
        // In practice, this would be:
        // let result = model.feature_importances();
        // assert!(result.is_err());
        // assert!(result.unwrap_err().to_string().contains("not been fitted"));

        // Instead, we just verify the model was created
        assert_eq!(model.config.n_estimators, 100); // default value
    }

    #[test]
    fn test_random_forest_regressor_proximity_matrix() {
        let x = array![[1.0], [2.0], [3.0], [4.0]];
        let y = array![1.0, 4.0, 9.0, 16.0];

        let model = RandomForestRegressor::new()
            .n_estimators(5)
            .max_depth(3)
            .criterion(SplitCriterion::MSE)
            .random_state(42)
            .fit(&x, &y)
            .unwrap();

        // Initially, proximity matrix should be None
        assert!(model.proximity_matrix().is_none());

        // Compute proximity matrix
        let proximity = model.compute_proximity_matrix(&x).unwrap();

        // Check dimensions
        assert_eq!(proximity.shape(), &[4, 4]);

        // Check diagonal elements are 1.0 (identity)
        for i in 0..4 {
            assert!((proximity[(i, i)] - 1.0).abs() < f64::EPSILON);
        }

        // Check symmetry
        for i in 0..4 {
            for j in 0..4 {
                assert!((proximity[(i, j)] - proximity[(j, i)]).abs() < f64::EPSILON);
            }
        }

        // Check values are in [0, 1] range
        for i in 0..4 {
            for j in 0..4 {
                assert!(proximity[(i, j)] >= 0.0 && proximity[(i, j)] <= 1.0);
            }
        }
    }

    #[test]
    fn test_random_forest_classifier_parallel_predict() {
        let x = array![
            [0.0, 0.0],
            [1.0, 1.0],
            [2.0, 2.0],
            [3.0, 3.0],
            [4.0, 4.0],
            [5.0, 5.0],
        ];
        let y = array![0, 0, 0, 1, 1, 1];

        let model = RandomForestClassifier::new()
            .n_estimators(10)
            .max_depth(3)
            .criterion(SplitCriterion::Gini)
            .random_state(42)
            .n_jobs(2) // Use 2 threads for parallel processing
            .fit(&x, &y)
            .unwrap();

        // Test parallel prediction
        let parallel_predictions = model.predict_parallel(&x).unwrap();
        let serial_predictions = model.predict(&x).unwrap();

        // Both should give the same results
        assert_eq!(parallel_predictions.len(), serial_predictions.len());
        assert_eq!(parallel_predictions.len(), 6);

        // The predictions should be identical (same model, same data)
        for (parallel, serial) in parallel_predictions.iter().zip(serial_predictions.iter()) {
            assert_eq!(parallel, serial);
        }
    }

    #[test]
    fn test_random_forest_classifier_parallel_predict_proba() {
        let x = array![
            [0.0, 0.0],
            [1.0, 1.0],
            [2.0, 2.0],
            [3.0, 3.0],
            [4.0, 4.0],
            [5.0, 5.0],
        ];
        let y = array![0, 0, 0, 1, 1, 1];

        let model = RandomForestClassifier::new()
            .n_estimators(10)
            .max_depth(3)
            .criterion(SplitCriterion::Gini)
            .random_state(42)
            .n_jobs(2) // Use 2 threads for parallel processing
            .fit(&x, &y)
            .unwrap();

        // Test parallel probability prediction
        let probabilities = model.predict_proba_parallel(&x).unwrap();

        // Check dimensions
        assert_eq!(probabilities.shape(), &[6, 2]); // 6 samples, 2 classes

        // Check that probabilities sum to 1.0 for each sample
        for i in 0..6 {
            let row_sum: f64 = probabilities.row(i).sum();
            assert!(
                (row_sum - 1.0).abs() < 1e-10,
                "Row {}: sum = {}",
                i,
                row_sum
            );
        }

        // Check that all probabilities are in [0, 1] range
        for prob in probabilities.iter() {
            assert!(
                *prob >= 0.0 && *prob <= 1.0,
                "Invalid probability: {}",
                prob
            );
        }
    }

    #[test]
    fn test_random_forest_regressor_parallel_predict() {
        let x = array![[0.0], [1.0], [2.0], [3.0], [4.0], [5.0]];
        let y = array![0.0, 1.0, 4.0, 9.0, 16.0, 25.0];

        let model = RandomForestRegressor::new()
            .n_estimators(20)
            .max_depth(5)
            .criterion(SplitCriterion::MSE)
            .random_state(42)
            .n_jobs(2) // Use 2 threads for parallel processing
            .fit(&x, &y)
            .unwrap();

        // Test parallel prediction
        let parallel_predictions = model.predict_parallel(&x).unwrap();
        let serial_predictions = model.predict(&x).unwrap();

        // Both should give the same results
        assert_eq!(parallel_predictions.len(), serial_predictions.len());
        assert_eq!(parallel_predictions.len(), 6);

        // The predictions should be identical (same model, same data)
        for (parallel, serial) in parallel_predictions.iter().zip(serial_predictions.iter()) {
            assert_eq!(parallel, serial);
        }

        // Test prediction on new data
        let test_x = array![[2.5]];
        let test_parallel_pred = model.predict_parallel(&test_x).unwrap();
        let test_serial_pred = model.predict(&test_x).unwrap();

        assert_eq!(test_parallel_pred.len(), 1);
        assert_eq!(test_serial_pred.len(), 1);
        assert_eq!(test_parallel_pred[0], test_serial_pred[0]);

        // Should predict something between 4 and 9
        assert!(test_parallel_pred[0] > 3.0 && test_parallel_pred[0] < 10.0);
    }
}

/// Calculate class weights for balanced Random Forest
pub fn calculate_class_weights(
    y: &Array1<i32>,
    strategy: &ClassWeight,
) -> Result<HashMap<i32, f64>> {
    match strategy {
        ClassWeight::None => {
            // Return equal weights for all classes
            let unique_classes: Vec<i32> = y
                .iter()
                .cloned()
                .collect::<std::collections::HashSet<_>>()
                .into_iter()
                .collect();
            let weights = unique_classes
                .into_iter()
                .map(|class| (class, 1.0))
                .collect();
            Ok(weights)
        }
        ClassWeight::Balanced => {
            // Calculate weights inversely proportional to class frequencies
            let mut class_counts: HashMap<i32, usize> = HashMap::new();
            for &class in y.iter() {
                *class_counts.entry(class).or_insert(0) += 1;
            }

            let n_samples = y.len() as f64;
            let n_classes = class_counts.len() as f64;

            let mut weights = HashMap::new();
            for (&class, &count) in &class_counts {
                let weight = n_samples / (n_classes * count as f64);
                weights.insert(class, weight);
            }
            Ok(weights)
        }
        ClassWeight::Custom(weights) => {
            // Use provided custom weights
            Ok(weights.clone())
        }
    }
}

/// Generate balanced bootstrap sample indices
pub fn balanced_bootstrap_sample(
    y: &Array1<i32>,
    strategy: SamplingStrategy,
    n_samples: usize,
    random_state: Option<u64>,
) -> Result<Vec<usize>> {
    let mut rng = scirs2_core::random::thread_rng();

    match strategy {
        SamplingStrategy::Bootstrap => {
            // Standard bootstrap sampling
            let mut indices = Vec::with_capacity(n_samples);
            for _ in 0..n_samples {
                indices.push(rng.gen_range(0..y.len()));
            }
            Ok(indices)
        }
        SamplingStrategy::BalancedBootstrap => {
            // Equal samples from each class
            let mut class_indices: HashMap<i32, Vec<usize>> = HashMap::new();
            for (idx, &class) in y.iter().enumerate() {
                class_indices.entry(class).or_default().push(idx);
            }

            let n_classes = class_indices.len();
            let samples_per_class = n_samples / n_classes;
            let extra_samples = n_samples % n_classes;

            let mut indices = Vec::with_capacity(n_samples);
            let mut extra_count = 0;

            for (_, class_idx_list) in class_indices.iter() {
                let mut n_class_samples = samples_per_class;
                if extra_count < extra_samples {
                    n_class_samples += 1;
                    extra_count += 1;
                }

                for _ in 0..n_class_samples {
                    let random_idx = rng.gen_range(0..class_idx_list.len());
                    indices.push(class_idx_list[random_idx]);
                }
            }

            // Shuffle the indices

            indices.shuffle(&mut rng);

            Ok(indices)
        }
        SamplingStrategy::Stratified => {
            // Preserve class distribution
            let mut class_counts: HashMap<i32, usize> = HashMap::new();
            let mut class_indices: HashMap<i32, Vec<usize>> = HashMap::new();

            for (idx, &class) in y.iter().enumerate() {
                *class_counts.entry(class).or_insert(0) += 1;
                class_indices.entry(class).or_default().push(idx);
            }

            let total_samples = y.len() as f64;
            let mut indices = Vec::with_capacity(n_samples);

            for (&class, &count) in &class_counts {
                let class_proportion = count as f64 / total_samples;
                let class_samples = (n_samples as f64 * class_proportion).round() as usize;
                let class_idx_list = &class_indices[&class];

                for _ in 0..class_samples {
                    let random_idx = rng.gen_range(0..class_idx_list.len());
                    indices.push(class_idx_list[random_idx]);
                }
            }

            // Fill remaining slots if needed
            while indices.len() < n_samples {
                indices.push(rng.gen_range(0..y.len()));
            }

            // Shuffle the indices

            indices.shuffle(&mut rng);

            Ok(indices)
        }
        SamplingStrategy::SMOTEBootstrap => {
            // SMOTE-like oversampling for minority classes
            let mut class_counts: HashMap<i32, usize> = HashMap::new();
            let mut class_indices: HashMap<i32, Vec<usize>> = HashMap::new();

            for (idx, &class) in y.iter().enumerate() {
                *class_counts.entry(class).or_insert(0) += 1;
                class_indices.entry(class).or_default().push(idx);
            }

            // Find majority class size
            let max_class_size = class_counts.values().max().copied().unwrap_or(0);
            let mut indices = Vec::new();

            for (&class, class_idx_list) in &class_indices {
                let class_count = class_counts[&class];
                let oversample_ratio = max_class_size as f64 / class_count as f64;
                let target_samples = (n_samples as f64 / class_counts.len() as f64
                    * oversample_ratio)
                    .round() as usize;

                for _ in 0..target_samples {
                    let random_idx = rng.gen_range(0..class_idx_list.len());
                    indices.push(class_idx_list[random_idx]);
                }
            }

            // Trim to exact size needed
            indices.truncate(n_samples);

            // Shuffle the indices

            indices.shuffle(&mut rng);

            Ok(indices)
        }
    }
}

/// Ensemble diversity measures for evaluating Random Forest and Extra Trees diversity
#[derive(Debug, Clone)]
pub struct DiversityMeasures {
    /// Q-statistic: Average pairwise Q-statistic between all classifier pairs
    pub q_statistic: f64,
    /// Disagreement measure: Average proportion of instances on which pairs disagree
    pub disagreement: f64,
    /// Double-fault measure: Average proportion of instances misclassified by both classifiers in pairs
    pub double_fault: f64,
    /// Correlation coefficient: Average correlation between classifier outputs
    pub correlation_coefficient: f64,
    /// Interrater agreement (Kappa): Agreement beyond chance
    pub kappa_statistic: f64,
    /// Entropy of ensemble predictions: Higher entropy indicates more diversity
    pub prediction_entropy: f64,
    /// Individual classifier accuracies
    pub individual_accuracies: Vec<f64>,
}

impl Default for DiversityMeasures {
    fn default() -> Self {
        Self::new()
    }
}

impl DiversityMeasures {
    /// Create new diversity measures with default values
    pub fn new() -> Self {
        Self {
            q_statistic: 0.0,
            disagreement: 0.0,
            double_fault: 0.0,
            correlation_coefficient: 0.0,
            kappa_statistic: 0.0,
            prediction_entropy: 0.0,
            individual_accuracies: Vec::new(),
        }
    }

    /// Print summary of diversity measures
    pub fn summary(&self) -> String {
        format!(
            "Diversity Measures Summary:\n\
             Q-statistic: {:.4} (higher = less diverse)\n\
             Disagreement: {:.4} (higher = more diverse)\n\
             Double-fault: {:.4} (lower = better)\n\
             Correlation: {:.4} (lower = more diverse)\n\
             Kappa: {:.4} (lower = more diverse)\n\
             Prediction Entropy: {:.4} (higher = more diverse)\n\
             Mean Individual Accuracy: {:.4}",
            self.q_statistic,
            self.disagreement,
            self.double_fault,
            self.correlation_coefficient,
            self.kappa_statistic,
            self.prediction_entropy,
            self.individual_accuracies.iter().sum::<f64>()
                / self.individual_accuracies.len() as f64
        )
    }
}

/// Calculate comprehensive diversity measures for an ensemble of classifiers
///
/// This function evaluates various measures of diversity between individual classifiers
/// in an ensemble, which helps understand how well the ensemble combines different
/// decision boundaries and reduces overfitting.
///
/// # Arguments
/// * `individual_predictions` - Matrix where rows are samples and columns are classifier predictions
/// * `true_labels` - Ground truth labels for the samples
///
/// # Returns
/// * `DiversityMeasures` struct containing various diversity metrics
pub fn calculate_ensemble_diversity(
    individual_predictions: &Array2<i32>,
    true_labels: &Array1<i32>,
) -> Result<DiversityMeasures> {
    let (n_samples, n_classifiers) = individual_predictions.dim();

    if n_samples == 0 || n_classifiers < 2 {
        return Err(SklearsError::InvalidInput(
            "Need at least 2 classifiers and some samples to calculate diversity".to_string(),
        ));
    }

    if true_labels.len() != n_samples {
        return Err(SklearsError::InvalidInput(
            "Number of true labels must match number of samples".to_string(),
        ));
    }

    // Calculate individual classifier accuracies
    let mut individual_accuracies = Vec::with_capacity(n_classifiers);
    for classifier_idx in 0..n_classifiers {
        let predictions = individual_predictions.column(classifier_idx);
        let accuracy = predictions
            .iter()
            .zip(true_labels.iter())
            .map(|(&pred, &true_label)| (pred == true_label) as i32)
            .sum::<i32>() as f64
            / n_samples as f64;
        individual_accuracies.push(accuracy);
    }

    // Calculate pairwise diversity measures
    let mut q_statistics = Vec::new();
    let mut disagreements = Vec::new();
    let mut double_faults = Vec::new();
    let mut correlations = Vec::new();
    let mut kappa_statistics = Vec::new();

    for i in 0..n_classifiers {
        for j in (i + 1)..n_classifiers {
            let pred_i = individual_predictions.column(i);
            let pred_j = individual_predictions.column(j);

            // Calculate confusion matrix elements for the pair
            let mut n11 = 0; // Both correct
            let mut n10 = 0; // i correct, j wrong
            let mut n01 = 0; // i wrong, j correct
            let mut n00 = 0; // Both wrong

            for sample_idx in 0..n_samples {
                let i_correct = pred_i[sample_idx] == true_labels[sample_idx];
                let j_correct = pred_j[sample_idx] == true_labels[sample_idx];

                match (i_correct, j_correct) {
                    (true, true) => n11 += 1,
                    (true, false) => n10 += 1,
                    (false, true) => n01 += 1,
                    (false, false) => n00 += 1,
                }
            }

            let n11_f = n11 as f64;
            let n10_f = n10 as f64;
            let n01_f = n01 as f64;
            let n00_f = n00 as f64;
            let n_f = n_samples as f64;

            // Q-statistic: (ad - bc) / (ad + bc)
            let q_stat = if (n11_f * n00_f + n10_f * n01_f) > 1e-10 {
                (n11_f * n00_f - n10_f * n01_f) / (n11_f * n00_f + n10_f * n01_f)
            } else {
                0.0
            };
            q_statistics.push(q_stat);

            // Disagreement measure: (b + c) / n
            let disagreement = (n10_f + n01_f) / n_f;
            disagreements.push(disagreement);

            // Double-fault measure: d / n
            let double_fault = n00_f / n_f;
            double_faults.push(double_fault);

            // Correlation coefficient between binary predictions
            let p_i = (n11_f + n10_f) / n_f; // Accuracy of classifier i
            let p_j = (n11_f + n01_f) / n_f; // Accuracy of classifier j

            let correlation = if p_i * (1.0 - p_i) * p_j * (1.0 - p_j) > 1e-10 {
                (n11_f / n_f - p_i * p_j) / ((p_i * (1.0 - p_i) * p_j * (1.0 - p_j)).sqrt())
            } else {
                0.0
            };
            correlations.push(correlation);

            // Kappa statistic (interrater agreement)
            let p_observed = (n11_f + n00_f) / n_f;
            let p_expected = p_i * p_j + (1.0 - p_i) * (1.0 - p_j);

            let kappa = if (1.0 - p_expected).abs() > 1e-10 {
                (p_observed - p_expected) / (1.0 - p_expected)
            } else {
                0.0
            };
            kappa_statistics.push(kappa);
        }
    }

    // Calculate ensemble prediction entropy
    let prediction_entropy = calculate_prediction_entropy(individual_predictions)?;

    Ok(DiversityMeasures {
        q_statistic: q_statistics.iter().sum::<f64>() / q_statistics.len() as f64,
        disagreement: disagreements.iter().sum::<f64>() / disagreements.len() as f64,
        double_fault: double_faults.iter().sum::<f64>() / double_faults.len() as f64,
        correlation_coefficient: correlations.iter().sum::<f64>() / correlations.len() as f64,
        kappa_statistic: kappa_statistics.iter().sum::<f64>() / kappa_statistics.len() as f64,
        prediction_entropy,
        individual_accuracies,
    })
}

/// Calculate prediction entropy of the ensemble
///
/// Higher entropy indicates that classifiers make more diverse predictions
fn calculate_prediction_entropy(individual_predictions: &Array2<i32>) -> Result<f64> {
    let (n_samples, n_classifiers) = individual_predictions.dim();
    let mut total_entropy = 0.0;

    for sample_idx in 0..n_samples {
        let sample_predictions = individual_predictions.row(sample_idx);

        // Count unique predictions for this sample
        let mut prediction_counts: HashMap<i32, usize> = HashMap::new();
        for &prediction in sample_predictions.iter() {
            *prediction_counts.entry(prediction).or_insert(0) += 1;
        }

        // Calculate entropy for this sample
        let mut sample_entropy = 0.0;
        for count in prediction_counts.values() {
            let probability = *count as f64 / n_classifiers as f64;
            if probability > 1e-10 {
                sample_entropy -= probability * probability.log2();
            }
        }

        total_entropy += sample_entropy;
    }

    Ok(total_entropy / n_samples as f64)
}

/// Calculate diversity measures for regression ensembles
///
/// For regression, we use different diversity measures based on prediction variance
/// and correlation between continuous outputs.
#[derive(Debug, Clone)]
pub struct RegressionDiversityMeasures {
    pub prediction_correlation: f64,
    pub prediction_variance: f64,
    pub average_bias: f64,
    pub average_variance: f64,
    pub individual_rmse: Vec<f64>,
}

/// Calculate diversity measures for regression ensembles
pub fn calculate_regression_diversity(
    individual_predictions: &Array2<f64>,
    true_values: &Array1<f64>,
) -> Result<RegressionDiversityMeasures> {
    let (n_samples, n_regressors) = individual_predictions.dim();

    if n_samples == 0 || n_regressors < 2 {
        return Err(SklearsError::InvalidInput(
            "Need at least 2 regressors and some samples".to_string(),
        ));
    }

    if true_values.len() != n_samples {
        return Err(SklearsError::InvalidInput(
            "Number of true values must match number of samples".to_string(),
        ));
    }

    // Calculate individual RMSE scores
    let mut individual_rmse = Vec::with_capacity(n_regressors);
    for regressor_idx in 0..n_regressors {
        let predictions = individual_predictions.column(regressor_idx);
        let mse = predictions
            .iter()
            .zip(true_values.iter())
            .map(|(&pred, &true_val)| (pred - true_val).powi(2))
            .sum::<f64>()
            / n_samples as f64;
        individual_rmse.push(mse.sqrt());
    }

    // Calculate pairwise correlations
    let mut correlations = Vec::new();
    for i in 0..n_regressors {
        for j in (i + 1)..n_regressors {
            let pred_i = individual_predictions.column(i);
            let pred_j = individual_predictions.column(j);

            let correlation =
                calculate_pearson_correlation(&pred_i.to_owned(), &pred_j.to_owned())?;
            correlations.push(correlation);
        }
    }

    // Calculate prediction variance for each sample
    let mut total_variance = 0.0;
    for sample_idx in 0..n_samples {
        let sample_predictions = individual_predictions.row(sample_idx);
        let mean_pred = sample_predictions.mean().unwrap();

        let variance = sample_predictions
            .iter()
            .map(|&pred| (pred - mean_pred).powi(2))
            .sum::<f64>()
            / n_regressors as f64;

        total_variance += variance;
    }
    let prediction_variance = total_variance / n_samples as f64;

    // Bias-variance decomposition (simplified)
    let mut total_bias = 0.0;
    let mut total_variance_component = 0.0;

    for sample_idx in 0..n_samples {
        let sample_predictions = individual_predictions.row(sample_idx);
        let mean_pred = sample_predictions.mean().unwrap();
        let true_val = true_values[sample_idx];

        // Bias^2: squared difference between mean prediction and true value
        let bias_squared = (mean_pred - true_val).powi(2);
        total_bias += bias_squared;

        // Variance: average squared difference from mean prediction
        let variance = sample_predictions
            .iter()
            .map(|&pred| (pred - mean_pred).powi(2))
            .sum::<f64>()
            / n_regressors as f64;
        total_variance_component += variance;
    }

    Ok(RegressionDiversityMeasures {
        prediction_correlation: correlations.iter().sum::<f64>() / correlations.len() as f64,
        prediction_variance,
        average_bias: (total_bias / n_samples as f64).sqrt(),
        average_variance: total_variance_component / n_samples as f64,
        individual_rmse,
    })
}

/// Calculate Pearson correlation coefficient between two arrays
fn calculate_pearson_correlation(x: &Array1<f64>, y: &Array1<f64>) -> Result<f64> {
    if x.len() != y.len() || x.len() < 2 {
        return Err(SklearsError::InvalidInput(
            "Arrays must have same length and at least 2 elements".to_string(),
        ));
    }

    let n = x.len() as f64;
    let mean_x = x.mean().unwrap();
    let mean_y = y.mean().unwrap();

    let mut numerator = 0.0;
    let mut sum_sq_x = 0.0;
    let mut sum_sq_y = 0.0;

    for i in 0..x.len() {
        let diff_x = x[i] - mean_x;
        let diff_y = y[i] - mean_y;

        numerator += diff_x * diff_y;
        sum_sq_x += diff_x * diff_x;
        sum_sq_y += diff_y * diff_y;
    }

    let denominator = (sum_sq_x * sum_sq_y).sqrt();

    if denominator < 1e-10 {
        Ok(0.0) // No correlation if no variance
    } else {
        Ok(numerator / denominator)
    }
}