ferrolearn-tree 0.5.0

//! Extremely randomized tree classifiers and regressors.
//!
//! This module provides [`ExtraTreeClassifier`] and [`ExtraTreeRegressor`],
//! which are variants of decision trees where split thresholds are chosen
//! randomly rather than via exhaustive search. For each candidate feature,
//! a random threshold is drawn uniformly between the feature's minimum and
//! maximum values in the current node.
//!
//! # Examples
//!
//! ```
//! use ferrolearn_tree::ExtraTreeClassifier;
//! use ferrolearn_core::{Fit, Predict};
//! use ndarray::{array, Array1, Array2};
//!
//! let x = Array2::from_shape_vec((6, 2), vec![
//!     1.0, 2.0,  2.0, 3.0,  3.0, 3.0,
//!     5.0, 6.0,  6.0, 7.0,  7.0, 8.0,
//! ]).unwrap();
//! let y = array![0, 0, 0, 1, 1, 1];
//!
//! let model = ExtraTreeClassifier::<f64>::new()
//!     .with_random_state(42);
//! let fitted = model.fit(&x, &y).unwrap();
//! let preds = fitted.predict(&x).unwrap();
//! ```
//!
//! ## REQ status
//!
//! Binary (R-DEFER-2): SHIPPED = impl + non-test consumer + tests + green
//! verification; NOT-STARTED = open blocker `#`. `ExtraTreeClassifier`/
//! `ExtraTreeRegressor` are re-exported at the crate root + registered as PyO3
//! `RsExtraTreeClassifier` (non-test consumers). ExtraTree is RNG-driven (random
//! split thresholds): exact node-for-node parity with sklearn at a given
//! `random_state` is the DOCUMENTED numpy-RNG boundary (#668/#670, same class as
//! SGD shuffle / libsvm CV) — verification pins the DETERMINISTIC contract +
//! ferrolearn-internal reproducibility. Pins in `tests/divergence_extra_tree.rs`.
//! See `.design/tree/extra_tree.md`.
//!
//! | REQ | Status | Evidence |
//! |---|---|---|
//! | REQ-1 (criteria gini/entropy/log_loss + mse/friedman_mse/absolute_error/poisson) | SHIPPED | classifier `compute_impurity` covers Gini/Entropy/LogLoss; `ExtraTreeRegressor` threads `criterion` through `RegressionData` → `regression_leaf_value` (median for AbsoluteError) + per-criterion split score in `find_random_regression_split` (mirrors `decision_tree.rs`); Poisson y>0 guard. Pinned by `pin3_reg_absolute_error_differs_from_mse_same_seed` + `pin3_reg_absolute_error_yields_median_leaf`. |
//! | REQ-2 (random-split FEATURE_THRESHOLD band + threshold==max guard) | NOT-STARTED | open prereq blocker #682. `find_random_*_split` use exact `feat_min >= feat_max` not sklearn's `1e-7` band (`_splitter.pyx:773`); no `threshold==max → min` guard (`:791`). |
//! | REQ-3 (param surface & defaults) | SHIPPED | clf `max_features=Sqrt`/`criterion=Gini`, reg `max_features=All`/`criterion=Mse`, both `splitter='random'`, `random_state=None` — match `_classes.py:1564`/`:1838` (live-verified). Pinned by `pin1_clf_defaults_match_sklearn`/`pin1_reg_defaults_match_sklearn`. `min_weight_fraction_leaf`/`max_leaf_nodes`/`min_impurity_decrease`/`ccp_alpha`/`class_weight` absent (#684). |
//! | REQ-4 (max_features resolution) | NOT-STARTED | open prereq blocker #685. `resolve_max_features` uses `ceil` vs sklearn's `int(...)` (off-by-one at e.g. n=2); no `max_features_` fitted attr (#686). |
//! | REQ-5 (fitted attributes) | SHIPPED | `fn feature_importances` (`HasFeatureImportances`, consumed by `extra_trees_ensemble.rs` + PyO3), `fn classes`/`n_classes`, `fn nodes`, `fn n_features`. `feature_importances_` exact value is RNG-dependent (documented boundary); `max_features_` NOT-STARTED (#686). |
//! | REQ-6 (predict / predict_proba) | SHIPPED | `fn predict`/`fn predict_proba`/`fn predict_log_proba` (consumed by `RsExtraTreeClassifier` + pipeline adapter); rows sum to 1. Exact leaf-frequency values RNG-dependent (boundary); multi-output NOT-STARTED (#689). |
//! | REQ-7 (random_state determinism) | SHIPPED | `random_state: Option<u64>` seeds `StdRng`; same seed ⇒ identical tree (`pin2_clf_random_state_reproducible`). Exact numpy-MT19937 cross-impl parity is the DOCUMENTED RNG boundary (#668/#670). |
//! | REQ-8 (ferray substrate) | NOT-STARTED | open prereq blocker #690. Imports `ndarray` + `rand::StdRng`, not `ferray-core`/`ferray::random` (R-SUBSTRATE). |
//! | REQ-9 (reject non-finite input) | SHIPPED | `fn reject_non_finite` called at the top of BOTH `ExtraTreeClassifier::fit` and `ExtraTreeRegressor::fit` (+ a float-`y` finite check in the regressor) rejects NaN AND infinity with `FerroError::InvalidParameter{name:"X"/"y"}`. UNLIKE the CART `DecisionTree*` base (which accepts NaN via `force_all_finite=False`, `_classes.py:248-250`), `ExtraTree*` use `splitter='random'` ⇒ `allow_nan=False` (`_classes.py:1085-1090`/`:1416-1421`) ⇒ `_support_missing_values` False ⇒ `assert_all_finite(X)` (`_classes.py:213-214`) raises `ValueError` (`validation.py:147-154`). Consumers: the existing `fit` entries (re-exported at the crate root + `RsExtraTreeClassifier` PyO3 reg). Pinned by `divergence_nan_reject_set.rs` (`divergence_extra_tree_classifier_nan_not_rejected`, `divergence_extra_tree_regressor_nan_not_rejected`) — live sklearn 1.5.2 raises, ferrolearn now `Err`. Finite input byte-identical (the module's oracle pins stay green). |

use ferrolearn_core::error::FerroError;
use ferrolearn_core::introspection::{HasClasses, HasFeatureImportances};
use ferrolearn_core::pipeline::{FittedPipelineEstimator, PipelineEstimator};
use ferrolearn_core::traits::{Fit, Predict};
use ndarray::{Array1, Array2};
use num_traits::{Float, FromPrimitive, ToPrimitive};
use rand::SeedableRng;
use rand::rngs::StdRng;
use rand::seq::index::sample as rand_sample_indices;
use serde::{Deserialize, Serialize};

use crate::decision_tree::{
    ClassificationCriterion, Node, RegressionCriterion, TreeParams, compute_feature_importances,
    traverse,
};
use crate::random_forest::MaxFeatures;

// ---------------------------------------------------------------------------
// Internal data structs for extra-tree building
// ---------------------------------------------------------------------------

/// Data references for classification extra-tree building.
struct ClassificationData<'a, F> {
    x: &'a Array2<F>,
    y: &'a [usize],
    n_classes: usize,
    feature_indices: Option<&'a [usize]>,
    criterion: ClassificationCriterion,
}

/// Data references for regression extra-tree building.
struct RegressionData<'a, F> {
    x: &'a Array2<F>,
    y: &'a Array1<F>,
    feature_indices: Option<&'a [usize]>,
    criterion: RegressionCriterion,
}

// ---------------------------------------------------------------------------
// ExtraTreeClassifier
// ---------------------------------------------------------------------------

/// Extremely randomized tree classifier.
///
/// Like a [`DecisionTreeClassifier`](crate::DecisionTreeClassifier), but split
/// thresholds are chosen randomly rather than via exhaustive search. For each
/// candidate feature, a random threshold is drawn uniformly between the
/// feature's minimum and maximum values in the current node.
///
/// # Type Parameters
///
/// - `F`: The floating-point type (`f32` or `f64`).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExtraTreeClassifier<F> {
    /// Maximum depth of the tree. `None` means unlimited.
    pub max_depth: Option<usize>,
    /// Minimum number of samples required to split an internal node.
    pub min_samples_split: usize,
    /// Minimum number of samples required in a leaf node.
    pub min_samples_leaf: usize,
    /// Strategy for the number of features considered at each split.
    pub max_features: MaxFeatures,
    /// Splitting criterion.
    pub criterion: ClassificationCriterion,
    /// Random seed for reproducibility. `None` means non-deterministic.
    pub random_state: Option<u64>,
    _marker: std::marker::PhantomData<F>,
}

impl<F: Float> ExtraTreeClassifier<F> {
    /// Create a new `ExtraTreeClassifier` with default settings.
    ///
    /// Defaults: `max_depth = None`, `min_samples_split = 2`,
    /// `min_samples_leaf = 1`, `max_features = Sqrt`,
    /// `criterion = Gini`, `random_state = None`.
    #[must_use]
    pub fn new() -> Self {
        Self {
            max_depth: None,
            min_samples_split: 2,
            min_samples_leaf: 1,
            max_features: MaxFeatures::Sqrt,
            criterion: ClassificationCriterion::Gini,
            random_state: None,
            _marker: std::marker::PhantomData,
        }
    }

    /// Set the maximum tree depth.
    #[must_use]
    pub fn with_max_depth(mut self, max_depth: Option<usize>) -> Self {
        self.max_depth = max_depth;
        self
    }

    /// Set the minimum number of samples required to split a node.
    #[must_use]
    pub fn with_min_samples_split(mut self, min_samples_split: usize) -> Self {
        self.min_samples_split = min_samples_split;
        self
    }

    /// Set the minimum number of samples required in a leaf node.
    #[must_use]
    pub fn with_min_samples_leaf(mut self, min_samples_leaf: usize) -> Self {
        self.min_samples_leaf = min_samples_leaf;
        self
    }

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

    /// Set the splitting criterion.
    #[must_use]
    pub fn with_criterion(mut self, criterion: ClassificationCriterion) -> Self {
        self.criterion = criterion;
        self
    }

    /// Set the random seed for reproducibility.
    #[must_use]
    pub fn with_random_state(mut self, seed: u64) -> Self {
        self.random_state = Some(seed);
        self
    }
}

impl<F: Float> Default for ExtraTreeClassifier<F> {
    fn default() -> Self {
        Self::new()
    }
}

// ---------------------------------------------------------------------------
// FittedExtraTreeClassifier
// ---------------------------------------------------------------------------

/// A fitted extremely randomized tree classifier.
///
/// Stores the learned tree as a flat `Vec<Node<F>>` for cache-friendly traversal.
/// Implements [`Predict`] for generating class predictions and
/// [`HasFeatureImportances`] for inspecting per-feature importance scores.
#[derive(Debug, Clone)]
pub struct FittedExtraTreeClassifier<F> {
    /// Flat node storage; index 0 is the root.
    nodes: Vec<Node<F>>,
    /// Sorted unique class labels observed during training.
    classes: Vec<usize>,
    /// Number of features the model was trained on.
    n_features: usize,
    /// Per-feature importance scores (normalised to sum to 1).
    feature_importances: Array1<F>,
}

impl<F: Float + Send + Sync + 'static> FittedExtraTreeClassifier<F> {
    /// Returns a reference to the flat node storage of the tree.
    #[must_use]
    pub fn nodes(&self) -> &[Node<F>] {
        &self.nodes
    }

    /// Returns the number of features the model was trained on.
    #[must_use]
    pub fn n_features(&self) -> usize {
        self.n_features
    }

    /// Predict class probabilities for each sample.
    ///
    /// Returns a 2-D array of shape `(n_samples, n_classes)`.
    ///
    /// # Errors
    ///
    /// Returns [`FerroError::ShapeMismatch`] if the number of features does
    /// not match the training data.
    pub fn predict_proba(&self, x: &Array2<F>) -> Result<Array2<F>, FerroError> {
        if x.ncols() != self.n_features {
            return Err(FerroError::ShapeMismatch {
                expected: vec![self.n_features],
                actual: vec![x.ncols()],
                context: "number of features must match fitted model".into(),
            });
        }
        let n_samples = x.nrows();
        let n_classes = self.classes.len();
        let mut proba = Array2::zeros((n_samples, n_classes));
        for i in 0..n_samples {
            let row = x.row(i);
            let leaf = traverse(&self.nodes, &row);
            if let Node::Leaf {
                class_distribution: Some(ref dist),
                ..
            } = self.nodes[leaf]
            {
                for (j, &p) in dist.iter().enumerate() {
                    proba[[i, j]] = p;
                }
            }
        }
        Ok(proba)
    }

    /// Mean accuracy on the given test data and labels.
    /// Equivalent to sklearn's `ClassifierMixin.score`.
    ///
    /// # Errors
    ///
    /// Returns [`FerroError::ShapeMismatch`] if `x.nrows() != y.len()` or
    /// the feature count does not match the training data.
    pub fn score(&self, x: &Array2<F>, y: &Array1<usize>) -> Result<F, FerroError> {
        if x.nrows() != y.len() {
            return Err(FerroError::ShapeMismatch {
                expected: vec![x.nrows()],
                actual: vec![y.len()],
                context: "y length must match number of samples in X".into(),
            });
        }
        let preds = self.predict(x)?;
        Ok(crate::mean_accuracy(&preds, y))
    }

    /// Element-wise log of [`predict_proba`](Self::predict_proba). Mirrors
    /// sklearn's `ClassifierMixin.predict_log_proba`.
    ///
    /// # Errors
    ///
    /// Forwards any error from [`predict_proba`](Self::predict_proba).
    pub fn predict_log_proba(&self, x: &Array2<F>) -> Result<Array2<F>, FerroError> {
        let proba = self.predict_proba(x)?;
        Ok(crate::log_proba(&proba))
    }
}

impl<F: Float + Send + Sync + 'static> Fit<Array2<F>, Array1<usize>> for ExtraTreeClassifier<F> {
    type Fitted = FittedExtraTreeClassifier<F>;
    type Error = FerroError;

    /// Fit the extra-tree classifier on the training data.
    ///
    /// # Errors
    ///
    /// Returns [`FerroError::ShapeMismatch`] if `x` and `y` have different
    /// numbers of samples.
    /// Returns [`FerroError::InsufficientSamples`] if there are no samples.
    /// Returns [`FerroError::InvalidParameter`] if hyperparameters are invalid.
    fn fit(
        &self,
        x: &Array2<F>,
        y: &Array1<usize>,
    ) -> Result<FittedExtraTreeClassifier<F>, FerroError> {
        let (n_samples, n_features) = x.dim();

        if n_samples != y.len() {
            return Err(FerroError::ShapeMismatch {
                expected: vec![n_samples],
                actual: vec![y.len()],
                context: "y length must match number of samples in X".into(),
            });
        }
        if n_samples == 0 {
            return Err(FerroError::InsufficientSamples {
                required: 1,
                actual: 0,
                context: "ExtraTreeClassifier requires at least one sample".into(),
            });
        }
        if self.min_samples_split < 2 {
            return Err(FerroError::InvalidParameter {
                name: "min_samples_split".into(),
                reason: "must be at least 2".into(),
            });
        }
        if self.min_samples_leaf < 1 {
            return Err(FerroError::InvalidParameter {
                name: "min_samples_leaf".into(),
                reason: "must be at least 1".into(),
            });
        }
        // Reject non-finite X up front (random splitter has no missing-value
        // support — `_classes.py:213-214`/`:1085-1090`).
        reject_non_finite(x)?;

        // Determine unique classes.
        let mut classes: Vec<usize> = y.iter().copied().collect();
        classes.sort_unstable();
        classes.dedup();
        let n_classes = classes.len();

        // Map class labels to indices 0..n_classes.
        let y_mapped: Vec<usize> = y
            .iter()
            .map(|&c| classes.iter().position(|&cl| cl == c).unwrap())
            .collect();

        let indices: Vec<usize> = (0..n_samples).collect();

        let max_features_n = resolve_max_features(self.max_features, n_features);

        let mut rng = if let Some(seed) = self.random_state {
            StdRng::seed_from_u64(seed)
        } else {
            StdRng::from_os_rng()
        };

        let data = ClassificationData {
            x,
            y: &y_mapped,
            n_classes,
            feature_indices: None,
            criterion: self.criterion,
        };
        let params = TreeParams {
            max_depth: self.max_depth,
            min_samples_split: self.min_samples_split,
            min_samples_leaf: self.min_samples_leaf,
        };

        let mut nodes: Vec<Node<F>> = Vec::new();
        build_extra_classification_tree(
            &data,
            &indices,
            &mut nodes,
            0,
            &params,
            n_features,
            max_features_n,
            &mut rng,
        );

        let feature_importances = compute_feature_importances(&nodes, n_features, n_samples);

        Ok(FittedExtraTreeClassifier {
            nodes,
            classes,
            n_features,
            feature_importances,
        })
    }
}

impl<F: Float + Send + Sync + 'static> Predict<Array2<F>> for FittedExtraTreeClassifier<F> {
    type Output = Array1<usize>;
    type Error = FerroError;

    /// Predict class labels for the given feature matrix.
    ///
    /// # Errors
    ///
    /// Returns [`FerroError::ShapeMismatch`] if the number of features does
    /// not match the fitted model.
    fn predict(&self, x: &Array2<F>) -> Result<Array1<usize>, FerroError> {
        if x.ncols() != self.n_features {
            return Err(FerroError::ShapeMismatch {
                expected: vec![self.n_features],
                actual: vec![x.ncols()],
                context: "number of features must match fitted model".into(),
            });
        }
        let n_samples = x.nrows();
        let mut predictions = Array1::zeros(n_samples);
        for i in 0..n_samples {
            let row = x.row(i);
            let leaf = traverse(&self.nodes, &row);
            if let Node::Leaf { value, .. } = self.nodes[leaf] {
                predictions[i] = float_to_usize(value);
            }
        }
        Ok(predictions)
    }
}

impl<F: Float + Send + Sync + 'static> HasFeatureImportances<F> for FittedExtraTreeClassifier<F> {
    fn feature_importances(&self) -> &Array1<F> {
        &self.feature_importances
    }
}

impl<F: Float + Send + Sync + 'static> HasClasses for FittedExtraTreeClassifier<F> {
    fn classes(&self) -> &[usize] {
        &self.classes
    }

    fn n_classes(&self) -> usize {
        self.classes.len()
    }
}

// Pipeline integration.
impl<F: Float + ToPrimitive + FromPrimitive + Send + Sync + 'static> PipelineEstimator<F>
    for ExtraTreeClassifier<F>
{
    fn fit_pipeline(
        &self,
        x: &Array2<F>,
        y: &Array1<F>,
    ) -> Result<Box<dyn FittedPipelineEstimator<F>>, FerroError> {
        let y_usize: Array1<usize> = y.mapv(|v| v.to_usize().unwrap_or(0));
        let fitted = self.fit(x, &y_usize)?;
        Ok(Box::new(FittedExtraTreeClassifierPipelineAdapter(fitted)))
    }
}

/// Pipeline adapter for `FittedExtraTreeClassifier<F>`.
struct FittedExtraTreeClassifierPipelineAdapter<F: Float + Send + Sync + 'static>(
    FittedExtraTreeClassifier<F>,
);

impl<F: Float + ToPrimitive + FromPrimitive + Send + Sync + 'static> FittedPipelineEstimator<F>
    for FittedExtraTreeClassifierPipelineAdapter<F>
{
    fn predict_pipeline(&self, x: &Array2<F>) -> Result<Array1<F>, FerroError> {
        let preds = self.0.predict(x)?;
        Ok(preds.mapv(|v| F::from_usize(v).unwrap_or_else(F::nan)))
    }
}

// ---------------------------------------------------------------------------
// ExtraTreeRegressor
// ---------------------------------------------------------------------------

/// Extremely randomized tree regressor.
///
/// Like a [`DecisionTreeRegressor`](crate::DecisionTreeRegressor), but split
/// thresholds are chosen randomly rather than via exhaustive search. For each
/// candidate feature, a random threshold is drawn uniformly between the
/// feature's minimum and maximum values in the current node.
///
/// # Type Parameters
///
/// - `F`: The floating-point type (`f32` or `f64`).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExtraTreeRegressor<F> {
    /// Maximum depth of the tree. `None` means unlimited.
    pub max_depth: Option<usize>,
    /// Minimum number of samples required to split an internal node.
    pub min_samples_split: usize,
    /// Minimum number of samples required in a leaf node.
    pub min_samples_leaf: usize,
    /// Strategy for the number of features considered at each split.
    pub max_features: MaxFeatures,
    /// Splitting criterion.
    pub criterion: RegressionCriterion,
    /// Random seed for reproducibility. `None` means non-deterministic.
    pub random_state: Option<u64>,
    _marker: std::marker::PhantomData<F>,
}

impl<F: Float> ExtraTreeRegressor<F> {
    /// Create a new `ExtraTreeRegressor` with default settings.
    ///
    /// Defaults: `max_depth = None`, `min_samples_split = 2`,
    /// `min_samples_leaf = 1`, `max_features = All`,
    /// `criterion = MSE`, `random_state = None`.
    #[must_use]
    pub fn new() -> Self {
        Self {
            max_depth: None,
            min_samples_split: 2,
            min_samples_leaf: 1,
            max_features: MaxFeatures::All,
            criterion: RegressionCriterion::Mse,
            random_state: None,
            _marker: std::marker::PhantomData,
        }
    }

    /// Set the maximum tree depth.
    #[must_use]
    pub fn with_max_depth(mut self, max_depth: Option<usize>) -> Self {
        self.max_depth = max_depth;
        self
    }

    /// Set the minimum number of samples required to split a node.
    #[must_use]
    pub fn with_min_samples_split(mut self, min_samples_split: usize) -> Self {
        self.min_samples_split = min_samples_split;
        self
    }

    /// Set the minimum number of samples required in a leaf node.
    #[must_use]
    pub fn with_min_samples_leaf(mut self, min_samples_leaf: usize) -> Self {
        self.min_samples_leaf = min_samples_leaf;
        self
    }

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

    /// Set the splitting criterion.
    #[must_use]
    pub fn with_criterion(mut self, criterion: RegressionCriterion) -> Self {
        self.criterion = criterion;
        self
    }

    /// Set the random seed for reproducibility.
    #[must_use]
    pub fn with_random_state(mut self, seed: u64) -> Self {
        self.random_state = Some(seed);
        self
    }
}

impl<F: Float> Default for ExtraTreeRegressor<F> {
    fn default() -> Self {
        Self::new()
    }
}

// ---------------------------------------------------------------------------
// FittedExtraTreeRegressor
// ---------------------------------------------------------------------------

/// A fitted extremely randomized tree regressor.
///
/// Stores the learned tree as a flat `Vec<Node<F>>` for cache-friendly traversal.
#[derive(Debug, Clone)]
pub struct FittedExtraTreeRegressor<F> {
    /// Flat node storage; index 0 is the root.
    nodes: Vec<Node<F>>,
    /// Number of features the model was trained on.
    n_features: usize,
    /// Per-feature importance scores (normalised to sum to 1).
    feature_importances: Array1<F>,
}

impl<F: Float + Send + Sync + 'static> FittedExtraTreeRegressor<F> {
    /// Returns a reference to the flat node storage of the tree.
    #[must_use]
    pub fn nodes(&self) -> &[Node<F>] {
        &self.nodes
    }

    /// Returns the number of features the model was trained on.
    #[must_use]
    pub fn n_features(&self) -> usize {
        self.n_features
    }

    /// R² coefficient of determination on the given test data.
    /// Equivalent to sklearn's `RegressorMixin.score`.
    ///
    /// # Errors
    ///
    /// Returns [`FerroError::ShapeMismatch`] if `x.nrows() != y.len()` or
    /// the feature count does not match the training data.
    pub fn score(&self, x: &Array2<F>, y: &Array1<F>) -> Result<F, FerroError> {
        if x.nrows() != y.len() {
            return Err(FerroError::ShapeMismatch {
                expected: vec![x.nrows()],
                actual: vec![y.len()],
                context: "y length must match number of samples in X".into(),
            });
        }
        let preds = self.predict(x)?;
        Ok(crate::r2_score(&preds, y))
    }
}

impl<F: Float + Send + Sync + 'static> Fit<Array2<F>, Array1<F>> for ExtraTreeRegressor<F> {
    type Fitted = FittedExtraTreeRegressor<F>;
    type Error = FerroError;

    /// Fit the extra-tree regressor on the training data.
    ///
    /// # Errors
    ///
    /// Returns [`FerroError::ShapeMismatch`] if `x` and `y` have different
    /// numbers of samples.
    /// Returns [`FerroError::InsufficientSamples`] if there are no samples.
    /// Returns [`FerroError::InvalidParameter`] if hyperparameters are invalid.
    fn fit(&self, x: &Array2<F>, y: &Array1<F>) -> Result<FittedExtraTreeRegressor<F>, FerroError> {
        let (n_samples, n_features) = x.dim();

        if n_samples != y.len() {
            return Err(FerroError::ShapeMismatch {
                expected: vec![n_samples],
                actual: vec![y.len()],
                context: "y length must match number of samples in X".into(),
            });
        }
        if n_samples == 0 {
            return Err(FerroError::InsufficientSamples {
                required: 1,
                actual: 0,
                context: "ExtraTreeRegressor requires at least one sample".into(),
            });
        }
        if self.min_samples_split < 2 {
            return Err(FerroError::InvalidParameter {
                name: "min_samples_split".into(),
                reason: "must be at least 2".into(),
            });
        }
        if self.min_samples_leaf < 1 {
            return Err(FerroError::InvalidParameter {
                name: "min_samples_leaf".into(),
                reason: "must be at least 1".into(),
            });
        }
        // Reject non-finite X (and the float target y) up front: the random
        // splitter has no missing-value support, and sklearn validates X/y
        // finiteness in `_validate_data` before any criterion-specific check
        // (`_classes.py:213-214`/`:1416-1421`).
        reject_non_finite(x)?;
        if y.iter().any(|v| !v.is_finite()) {
            return Err(FerroError::InvalidParameter {
                name: "y".into(),
                reason: "Input y contains NaN or infinity.".into(),
            });
        }

        // Poisson requires non-negative targets with a strictly positive sum,
        // mirroring sklearn's check (`_classes.py:267-277`) and the analogous
        // guard in `decision_tree.rs` `DecisionTreeRegressor::fit`.
        if self.criterion == RegressionCriterion::Poisson {
            if y.iter().any(|&v| v < F::zero()) {
                return Err(FerroError::InvalidParameter {
                    name: "y".into(),
                    reason: "Some value(s) of y are negative which is not allowed for Poisson \
                             regression."
                        .into(),
                });
            }
            let sum_y = y.iter().fold(F::zero(), |a, &b| a + b);
            if sum_y <= F::zero() {
                return Err(FerroError::InvalidParameter {
                    name: "y".into(),
                    reason: "Sum of y is not positive which is necessary for Poisson regression."
                        .into(),
                });
            }
        }

        let indices: Vec<usize> = (0..n_samples).collect();
        let max_features_n = resolve_max_features(self.max_features, n_features);

        let mut rng = if let Some(seed) = self.random_state {
            StdRng::seed_from_u64(seed)
        } else {
            StdRng::from_os_rng()
        };

        let data = RegressionData {
            x,
            y,
            feature_indices: None,
            criterion: self.criterion,
        };
        let params = TreeParams {
            max_depth: self.max_depth,
            min_samples_split: self.min_samples_split,
            min_samples_leaf: self.min_samples_leaf,
        };

        let mut nodes: Vec<Node<F>> = Vec::new();
        build_extra_regression_tree(
            &data,
            &indices,
            &mut nodes,
            0,
            &params,
            n_features,
            max_features_n,
            &mut rng,
        );

        let feature_importances = compute_feature_importances(&nodes, n_features, n_samples);

        Ok(FittedExtraTreeRegressor {
            nodes,
            n_features,
            feature_importances,
        })
    }
}

impl<F: Float + Send + Sync + 'static> Predict<Array2<F>> for FittedExtraTreeRegressor<F> {
    type Output = Array1<F>;
    type Error = FerroError;

    /// Predict target values for the given feature matrix.
    ///
    /// # Errors
    ///
    /// Returns [`FerroError::ShapeMismatch`] if the number of features does
    /// not match the fitted model.
    fn predict(&self, x: &Array2<F>) -> Result<Array1<F>, FerroError> {
        if x.ncols() != self.n_features {
            return Err(FerroError::ShapeMismatch {
                expected: vec![self.n_features],
                actual: vec![x.ncols()],
                context: "number of features must match fitted model".into(),
            });
        }
        let n_samples = x.nrows();
        let mut predictions = Array1::zeros(n_samples);
        for i in 0..n_samples {
            let row = x.row(i);
            let leaf = traverse(&self.nodes, &row);
            if let Node::Leaf { value, .. } = self.nodes[leaf] {
                predictions[i] = value;
            }
        }
        Ok(predictions)
    }
}

impl<F: Float + Send + Sync + 'static> HasFeatureImportances<F> for FittedExtraTreeRegressor<F> {
    fn feature_importances(&self) -> &Array1<F> {
        &self.feature_importances
    }
}

// Pipeline integration.
impl<F: Float + Send + Sync + 'static> PipelineEstimator<F> for ExtraTreeRegressor<F> {
    fn fit_pipeline(
        &self,
        x: &Array2<F>,
        y: &Array1<F>,
    ) -> Result<Box<dyn FittedPipelineEstimator<F>>, FerroError> {
        let fitted = self.fit(x, y)?;
        Ok(Box::new(fitted))
    }
}

impl<F: Float + Send + Sync + 'static> FittedPipelineEstimator<F> for FittedExtraTreeRegressor<F> {
    fn predict_pipeline(&self, x: &Array2<F>) -> Result<Array1<F>, FerroError> {
        self.predict(x)
    }
}

// ---------------------------------------------------------------------------
// Internal: helpers
// ---------------------------------------------------------------------------

/// Reject `X` containing any non-finite value (NaN or infinity).
///
/// Mirrors sklearn's up-front `_validate_data(..., force_all_finite=True)` for
/// the random-splitter path: `ExtraTree*` use `splitter='random'`, so
/// `allow_nan` is `False` (`sklearn/tree/_classes.py:1085-1090`/`:1416-1421`),
/// `_support_missing_values` returns `False`, and the data check calls
/// `assert_all_finite(X)` (`_classes.py:213-214`), which raises
/// `ValueError("Input X contains NaN.")` (`sklearn/utils/validation.py:147-154`).
/// Unlike the CART `DecisionTree*` base (which accepts NaN via
/// `force_all_finite=False`), the random splitter does NOT support missing
/// values, so NaN AND infinity are both rejected. Never panics (R-CODE-2).
fn reject_non_finite<F: Float>(x: &Array2<F>) -> Result<(), FerroError> {
    if x.iter().any(|v| !v.is_finite()) {
        return Err(FerroError::InvalidParameter {
            name: "X".into(),
            reason: "Input X contains NaN or infinity.".into(),
        });
    }
    Ok(())
}

/// Resolve the `MaxFeatures` strategy to a concrete number.
fn resolve_max_features(strategy: MaxFeatures, n_features: usize) -> usize {
    let result = match strategy {
        MaxFeatures::Sqrt => (n_features as f64).sqrt().ceil() as usize,
        MaxFeatures::Log2 => (n_features as f64).log2().ceil().max(1.0) as usize,
        MaxFeatures::All => n_features,
        MaxFeatures::Fixed(n) => n.min(n_features),
        MaxFeatures::Fraction(f) => ((n_features as f64) * f).ceil() as usize,
    };
    result.max(1).min(n_features)
}

/// Convert a `Float` value to `usize` (for class labels stored as floats).
fn float_to_usize<F: Float>(v: F) -> usize {
    v.to_f64().map_or(0, |f| f.round() as usize)
}

/// Generate a uniform random float in `[min_val, max_val]`.
fn random_threshold<F: Float>(rng: &mut StdRng, min_val: F, max_val: F) -> F {
    use rand::RngCore;
    // Generate a random f64 in [0, 1) and scale to [min_val, max_val].
    let u = (rng.next_u64() as f64) / (u64::MAX as f64);
    let range = max_val - min_val;
    min_val + F::from(u).unwrap() * range
}

/// Compute the Gini impurity for a set of class counts.
fn gini_impurity<F: Float>(class_counts: &[usize], total: usize) -> F {
    if total == 0 {
        return F::zero();
    }
    let total_f = F::from(total).unwrap();
    let mut impurity = F::one();
    for &count in class_counts {
        let p = F::from(count).unwrap() / total_f;
        impurity = impurity - p * p;
    }
    impurity
}

/// Compute the Shannon entropy for a set of class counts.
fn entropy_impurity<F: Float>(class_counts: &[usize], total: usize) -> F {
    if total == 0 {
        return F::zero();
    }
    let total_f = F::from(total).unwrap();
    let mut ent = F::zero();
    for &count in class_counts {
        if count > 0 {
            let p = F::from(count).unwrap() / total_f;
            ent = ent - p * p.ln();
        }
    }
    ent
}

/// Compute impurity for a given classification criterion.
fn compute_impurity<F: Float>(
    class_counts: &[usize],
    total: usize,
    criterion: ClassificationCriterion,
) -> F {
    match criterion {
        ClassificationCriterion::Gini => gini_impurity(class_counts, total),
        // `log_loss` is an alias of `entropy` in sklearn's CRITERIA_CLF
        // (`sklearn/tree/_classes.py`: both map to `_criterion.Entropy`), so
        // LogLoss uses the identical Shannon-entropy impurity — mirroring
        // `decision_tree.rs`'s `Entropy | LogLoss => entropy_impurity` arm.
        ClassificationCriterion::Entropy | ClassificationCriterion::LogLoss => {
            entropy_impurity(class_counts, total)
        }
    }
}

/// Create a classification leaf node and return its index.
fn make_classification_leaf<F: Float>(
    nodes: &mut Vec<Node<F>>,
    class_counts: &[usize],
    n_classes: usize,
    n_samples: usize,
) -> usize {
    let majority_class = class_counts
        .iter()
        .enumerate()
        .max_by_key(|&(_, &count)| count)
        .map_or(0, |(i, _)| i);

    let total_f = if n_samples > 0 {
        F::from(n_samples).unwrap()
    } else {
        F::one()
    };
    let distribution: Vec<F> = (0..n_classes)
        .map(|c| F::from(class_counts[c]).unwrap() / total_f)
        .collect();

    let idx = nodes.len();
    nodes.push(Node::Leaf {
        value: F::from(majority_class).unwrap(),
        class_distribution: Some(distribution),
        n_samples,
    });
    idx
}

/// Compute the mean of target values for the given indices.
fn mean_value<F: Float>(y: &Array1<F>, indices: &[usize]) -> F {
    if indices.is_empty() {
        return F::zero();
    }
    let sum: F = indices.iter().map(|&i| y[i]).fold(F::zero(), |a, b| a + b);
    sum / F::from(indices.len()).unwrap()
}

/// Compute the median of target values for the given indices.
///
/// Mirrors `decision_tree.rs`'s `median_value` / sklearn's
/// `WeightedMedianCalculator.get_median` (unweighted, `MAE.node_value`,
/// `_criterion.pyx:1419-1423`): for an even-length sample the median is the
/// average of the two middle (sorted) values. Uses a NaN-last total order via
/// `partial_cmp(..).unwrap_or(Equal)` so it never panics (R-CODE-2 / R-APG-1).
fn median_value<F: Float>(y: &Array1<F>, indices: &[usize]) -> F {
    let n = indices.len();
    if n == 0 {
        return F::zero();
    }
    let mut vals: Vec<F> = indices.iter().map(|&i| y[i]).collect();
    vals.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
    let mid = n / 2;
    if n % 2 == 1 {
        vals[mid]
    } else {
        (vals[mid - 1] + vals[mid]) / F::from(2.0).unwrap_or_else(F::one)
    }
}

/// Compute the mean absolute error of the given indices around their median,
/// `(1/n)·Σ|y_i − median|` (`_criterion.pyx:1450-1472`, `MAE.node_impurity`).
fn mae_for_indices<F: Float>(y: &Array1<F>, indices: &[usize]) -> F {
    let n = indices.len();
    if n == 0 {
        return F::zero();
    }
    let median = median_value(y, indices);
    let sum_abs: F = indices
        .iter()
        .map(|&i| (y[i] - median).abs())
        .fold(F::zero(), |a, b| a + b);
    sum_abs / F::from(n).unwrap_or_else(F::one)
}

/// Compute the half-Poisson deviance impurity of the given indices,
/// `(1/n)·Σ y_i·ln(y_i/mean)` with `0·ln0 = 0` (`_criterion.pyx:1671-1708`,
/// `Poisson.poisson_loss`). Returns `+∞` when the node's target sum is
/// non-positive, mirroring sklearn's `y_sum <= EPSILON ⇒ return INFINITY`
/// guard so such splits are never selected.
fn poisson_deviance_for_indices<F: Float>(y: &Array1<F>, indices: &[usize]) -> F {
    let n = indices.len();
    if n == 0 {
        return F::zero();
    }
    let n_f = F::from(n).unwrap_or_else(F::one);
    let sum: F = indices.iter().map(|&i| y[i]).fold(F::zero(), |a, b| a + b);
    if sum <= F::epsilon() {
        return F::infinity();
    }
    let mean = sum / n_f;
    let mut loss = F::zero();
    for &i in indices {
        let yi = y[i];
        // xlogy(y, y/mean): the `y == 0` term contributes 0 (`0·ln0 = 0`).
        if yi > F::zero() {
            loss = loss + yi * (yi / mean).ln();
        }
    }
    loss / n_f
}

/// Compute the leaf prediction value for a regression node under `criterion`.
///
/// Mirrors `decision_tree.rs`'s `regression_leaf_value` / `Criterion.node_value`
/// (`_criterion.pyx`): MSE / FriedmanMSE / Poisson predict the **mean**
/// (`_criterion.pyx:1052`), while MAE / `absolute_error` predicts the **median**
/// (`MAE.node_value`, `_criterion.pyx:1419-1423`).
fn regression_leaf_value<F: Float>(
    y: &Array1<F>,
    indices: &[usize],
    criterion: RegressionCriterion,
) -> F {
    match criterion {
        RegressionCriterion::AbsoluteError => median_value(y, indices),
        RegressionCriterion::Mse
        | RegressionCriterion::FriedmanMse
        | RegressionCriterion::Poisson => mean_value(y, indices),
    }
}

/// Compute the node impurity for a regression node under `criterion`
/// (`Criterion.node_impurity`), mirroring `decision_tree.rs`'s
/// `regression_node_impurity`.
fn regression_node_impurity<F: Float>(
    y: &Array1<F>,
    indices: &[usize],
    criterion: RegressionCriterion,
) -> F {
    match criterion {
        // FriedmanMSE shares MSE's node impurity (it only overrides the split
        // improvement); MSE is the variance around the mean (`_criterion.pyx:1094`).
        RegressionCriterion::Mse | RegressionCriterion::FriedmanMse => {
            if indices.is_empty() {
                return F::zero();
            }
            let mean = mean_value(y, indices);
            let sum_sq: F = indices
                .iter()
                .map(|&i| {
                    let diff = y[i] - mean;
                    diff * diff
                })
                .fold(F::zero(), |a, b| a + b);
            sum_sq / F::from(indices.len()).unwrap_or_else(F::one)
        }
        RegressionCriterion::AbsoluteError => mae_for_indices(y, indices),
        RegressionCriterion::Poisson => poisson_deviance_for_indices(y, indices),
    }
}

// ---------------------------------------------------------------------------
// Extra-tree classification building
// ---------------------------------------------------------------------------

/// Build an extra-tree classification tree recursively with random thresholds.
///
/// At each node, a random subset of features is considered, and for each feature
/// a random threshold is drawn uniformly between the feature's min and max in the
/// current node.
#[allow(clippy::too_many_arguments)]
fn build_extra_classification_tree<F: Float>(
    data: &ClassificationData<'_, F>,
    indices: &[usize],
    nodes: &mut Vec<Node<F>>,
    depth: usize,
    params: &TreeParams,
    n_features: usize,
    max_features_n: usize,
    rng: &mut StdRng,
) -> usize {
    let n = indices.len();

    let mut class_counts = vec![0usize; data.n_classes];
    for &i in indices {
        class_counts[data.y[i]] += 1;
    }

    let should_stop = n < params.min_samples_split
        || params.max_depth.is_some_and(|d| depth >= d)
        || class_counts.iter().filter(|&&c| c > 0).count() <= 1;

    if should_stop {
        return make_classification_leaf(nodes, &class_counts, data.n_classes, n);
    }

    let best = find_random_classification_split(
        data,
        indices,
        params.min_samples_leaf,
        n_features,
        max_features_n,
        rng,
    );

    if let Some((best_feature, best_threshold, best_impurity_decrease)) = best {
        let (left_indices, right_indices): (Vec<usize>, Vec<usize>) = indices
            .iter()
            .partition(|&&i| data.x[[i, best_feature]] <= best_threshold);

        // Ensure both children have at least min_samples_leaf.
        if left_indices.len() < params.min_samples_leaf
            || right_indices.len() < params.min_samples_leaf
        {
            return make_classification_leaf(nodes, &class_counts, data.n_classes, n);
        }

        let node_idx = nodes.len();
        nodes.push(Node::Leaf {
            value: F::zero(),
            class_distribution: None,
            n_samples: 0,
        }); // placeholder

        let left_idx = build_extra_classification_tree(
            data,
            &left_indices,
            nodes,
            depth + 1,
            params,
            n_features,
            max_features_n,
            rng,
        );
        let right_idx = build_extra_classification_tree(
            data,
            &right_indices,
            nodes,
            depth + 1,
            params,
            n_features,
            max_features_n,
            rng,
        );

        nodes[node_idx] = Node::Split {
            feature: best_feature,
            threshold: best_threshold,
            left: left_idx,
            right: right_idx,
            impurity_decrease: best_impurity_decrease,
            n_samples: n,
        };

        node_idx
    } else {
        make_classification_leaf(nodes, &class_counts, data.n_classes, n)
    }
}

/// Find the best random split for a classification node.
///
/// For each candidate feature (from a random subset), pick a random threshold
/// between min and max of that feature in the current node. Return the split
/// with the largest impurity decrease, or `None` if no valid split exists.
#[allow(clippy::too_many_arguments)]
fn find_random_classification_split<F: Float>(
    data: &ClassificationData<'_, F>,
    indices: &[usize],
    min_samples_leaf: usize,
    n_features: usize,
    max_features_n: usize,
    rng: &mut StdRng,
) -> Option<(usize, F, F)> {
    let n = indices.len();
    let n_f = F::from(n).unwrap();

    let mut parent_counts = vec![0usize; data.n_classes];
    for &i in indices {
        parent_counts[data.y[i]] += 1;
    }
    let parent_impurity = compute_impurity::<F>(&parent_counts, n, data.criterion);

    let mut best_score = F::neg_infinity();
    let mut best_feature = 0;
    let mut best_threshold = F::zero();

    // Select random feature subset.
    let feature_subset: Vec<usize> = if let Some(feat_indices) = data.feature_indices {
        // If a feature subset is provided externally, sample from it.
        let k = max_features_n.min(feat_indices.len());
        rand_sample_indices(rng, feat_indices.len(), k)
            .into_vec()
            .into_iter()
            .map(|i| feat_indices[i])
            .collect()
    } else {
        let k = max_features_n.min(n_features);
        rand_sample_indices(rng, n_features, k).into_vec()
    };

    for feat in feature_subset {
        // Find min and max of this feature in the current node.
        let mut feat_min = F::infinity();
        let mut feat_max = F::neg_infinity();
        for &i in indices {
            let val = data.x[[i, feat]];
            if val < feat_min {
                feat_min = val;
            }
            if val > feat_max {
                feat_max = val;
            }
        }

        // Skip constant features.
        if feat_min >= feat_max {
            continue;
        }

        // Draw a random threshold uniformly in (min, max).
        let threshold = random_threshold(rng, feat_min, feat_max);

        // Count left and right.
        let mut left_counts = vec![0usize; data.n_classes];
        let mut right_counts = vec![0usize; data.n_classes];
        let mut left_n = 0usize;

        for &i in indices {
            let cls = data.y[i];
            if data.x[[i, feat]] <= threshold {
                left_counts[cls] += 1;
                left_n += 1;
            } else {
                right_counts[cls] += 1;
            }
        }

        let right_n = n - left_n;
        if left_n < min_samples_leaf || right_n < min_samples_leaf {
            continue;
        }

        let left_impurity = compute_impurity::<F>(&left_counts, left_n, data.criterion);
        let right_impurity = compute_impurity::<F>(&right_counts, right_n, data.criterion);
        let left_weight = F::from(left_n).unwrap() / n_f;
        let right_weight = F::from(right_n).unwrap() / n_f;
        let weighted_child_impurity = left_weight * left_impurity + right_weight * right_impurity;
        let impurity_decrease = parent_impurity - weighted_child_impurity;

        if impurity_decrease > best_score {
            best_score = impurity_decrease;
            best_feature = feat;
            best_threshold = threshold;
        }
    }

    if best_score > F::zero() {
        Some((best_feature, best_threshold, best_score * n_f))
    } else {
        None
    }
}

// ---------------------------------------------------------------------------
// Extra-tree regression building
// ---------------------------------------------------------------------------

/// Build an extra-tree regression tree recursively with random thresholds.
#[allow(clippy::too_many_arguments)]
fn build_extra_regression_tree<F: Float>(
    data: &RegressionData<'_, F>,
    indices: &[usize],
    nodes: &mut Vec<Node<F>>,
    depth: usize,
    params: &TreeParams,
    n_features: usize,
    max_features_n: usize,
    rng: &mut StdRng,
) -> usize {
    let n = indices.len();
    // Leaf prediction depends on the criterion: median for `AbsoluteError`,
    // mean for MSE / FriedmanMSE / Poisson (mirrors `regression_leaf_value`).
    let leaf_value = regression_leaf_value(data.y, indices, data.criterion);

    let should_stop = n < params.min_samples_split || params.max_depth.is_some_and(|d| depth >= d);

    if should_stop {
        let idx = nodes.len();
        nodes.push(Node::Leaf {
            value: leaf_value,
            class_distribution: None,
            n_samples: n,
        });
        return idx;
    }

    // Pure-node short-circuit: the node's impurity (variance for MSE/FriedmanMSE,
    // L1 spread for AbsoluteError, deviance for Poisson) is essentially zero ⇒
    // no split can improve it, so emit a leaf. Uses the criterion's own node
    // impurity so the stop condition is criterion-consistent.
    let parent_impurity = regression_node_impurity(data.y, indices, data.criterion);

    if parent_impurity <= F::epsilon() {
        let idx = nodes.len();
        nodes.push(Node::Leaf {
            value: leaf_value,
            class_distribution: None,
            n_samples: n,
        });
        return idx;
    }

    let best = find_random_regression_split(
        data,
        indices,
        params.min_samples_leaf,
        n_features,
        max_features_n,
        rng,
    );

    if let Some((best_feature, best_threshold, best_impurity_decrease)) = best {
        let (left_indices, right_indices): (Vec<usize>, Vec<usize>) = indices
            .iter()
            .partition(|&&i| data.x[[i, best_feature]] <= best_threshold);

        // Ensure both children have at least min_samples_leaf.
        if left_indices.len() < params.min_samples_leaf
            || right_indices.len() < params.min_samples_leaf
        {
            let idx = nodes.len();
            nodes.push(Node::Leaf {
                value: leaf_value,
                class_distribution: None,
                n_samples: n,
            });
            return idx;
        }

        let node_idx = nodes.len();
        nodes.push(Node::Leaf {
            value: F::zero(),
            class_distribution: None,
            n_samples: 0,
        }); // placeholder

        let left_idx = build_extra_regression_tree(
            data,
            &left_indices,
            nodes,
            depth + 1,
            params,
            n_features,
            max_features_n,
            rng,
        );
        let right_idx = build_extra_regression_tree(
            data,
            &right_indices,
            nodes,
            depth + 1,
            params,
            n_features,
            max_features_n,
            rng,
        );

        nodes[node_idx] = Node::Split {
            feature: best_feature,
            threshold: best_threshold,
            left: left_idx,
            right: right_idx,
            impurity_decrease: best_impurity_decrease,
            n_samples: n,
        };

        node_idx
    } else {
        let idx = nodes.len();
        nodes.push(Node::Leaf {
            value: leaf_value,
            class_distribution: None,
            n_samples: n,
        });
        idx
    }
}

/// Find the best random split for a regression node.
///
/// For each candidate feature (from a random subset), pick a random threshold
/// between min and max of that feature in the current node. Return the split
/// with the largest MSE decrease, or `None` if no valid split exists.
#[allow(clippy::too_many_arguments)]
fn find_random_regression_split<F: Float>(
    data: &RegressionData<'_, F>,
    indices: &[usize],
    min_samples_leaf: usize,
    n_features: usize,
    max_features_n: usize,
    rng: &mut StdRng,
) -> Option<(usize, F, F)> {
    let n = indices.len();
    let n_f = F::from(n).unwrap();

    let parent_sum: F = indices
        .iter()
        .map(|&i| data.y[i])
        .fold(F::zero(), |a, b| a + b);
    let parent_sum_sq: F = indices
        .iter()
        .map(|&i| data.y[i] * data.y[i])
        .fold(F::zero(), |a, b| a + b);
    let parent_mse = parent_sum_sq / n_f - (parent_sum / n_f) * (parent_sum / n_f);
    // Parent impurity for the median-based (absolute_error) and Poisson
    // criteria; unused (cheap to leave) for the variance-based MSE/FriedmanMSE
    // paths, which score off `parent_mse` / the Friedman proxy instead. Mirrors
    // `decision_tree.rs`'s `find_best_regression_split`.
    let parent_impurity = regression_node_impurity(data.y, indices, data.criterion);

    let mut best_score = F::neg_infinity();
    let mut best_feature = 0;
    let mut best_threshold = F::zero();

    // Select random feature subset.
    let feature_subset: Vec<usize> = if let Some(feat_indices) = data.feature_indices {
        let k = max_features_n.min(feat_indices.len());
        rand_sample_indices(rng, feat_indices.len(), k)
            .into_vec()
            .into_iter()
            .map(|i| feat_indices[i])
            .collect()
    } else {
        let k = max_features_n.min(n_features);
        rand_sample_indices(rng, n_features, k).into_vec()
    };

    for feat in feature_subset {
        // Find min and max of this feature in the current node.
        let mut feat_min = F::infinity();
        let mut feat_max = F::neg_infinity();
        for &i in indices {
            let val = data.x[[i, feat]];
            if val < feat_min {
                feat_min = val;
            }
            if val > feat_max {
                feat_max = val;
            }
        }

        // Skip constant features.
        if feat_min >= feat_max {
            continue;
        }

        // Draw a random threshold uniformly in (min, max).
        let threshold = random_threshold(rng, feat_min, feat_max);

        // Compute left/right statistics. The running sums power the
        // byte-identical MSE path; the index slices power the median-based
        // (absolute_error) and Poisson criteria.
        let mut left_sum = F::zero();
        let mut left_sum_sq = F::zero();
        let mut left_indices: Vec<usize> = Vec::new();
        let mut right_indices: Vec<usize> = Vec::new();

        for &i in indices {
            if data.x[[i, feat]] <= threshold {
                let val = data.y[i];
                left_sum = left_sum + val;
                left_sum_sq = left_sum_sq + val * val;
                left_indices.push(i);
            } else {
                right_indices.push(i);
            }
        }

        let left_n = left_indices.len();
        let right_n = n - left_n;
        if left_n < min_samples_leaf || right_n < min_samples_leaf {
            continue;
        }

        let left_n_f = F::from(left_n).unwrap_or_else(F::one);
        let right_n_f = F::from(right_n).unwrap_or_else(F::one);

        // Per-criterion split score (higher = better). All four reduce to a
        // "parent impurity minus weighted child impurity" improvement so the
        // shared `> best_score` argmax + `> 0` accept gate is reused. Mirrors
        // `decision_tree.rs`'s `find_best_regression_split`.
        let score = match data.criterion {
            // squared_error: parent_mse − (n_L·mse_L + n_R·mse_R)/n
            // (`_criterion.pyx:1094`), kept byte-identical to the prior MSE path.
            RegressionCriterion::Mse => {
                let left_mean = left_sum / left_n_f;
                let left_mse = left_sum_sq / left_n_f - left_mean * left_mean;
                let right_sum = parent_sum - left_sum;
                let right_sum_sq = parent_sum_sq - left_sum_sq;
                let right_mean = right_sum / right_n_f;
                let right_mse = right_sum_sq / right_n_f - right_mean * right_mean;
                let weighted_child_mse = (left_n_f * left_mse + right_n_f * right_mse) / n_f;
                parent_mse - weighted_child_mse
            }
            // friedman_mse proxy (`FriedmanMSE.impurity_improvement`,
            // `_criterion.pyx:1557-1574`, n_outputs == 1):
            //   diff = n_R·sum_L − n_L·sum_R; improvement = diff²/(n_L·n_R·n_node).
            RegressionCriterion::FriedmanMse => {
                let right_sum = parent_sum - left_sum;
                let diff = right_n_f * left_sum - left_n_f * right_sum;
                diff * diff / (left_n_f * right_n_f * n_f)
            }
            // absolute_error: parent_mae − (n_L·mae_L + n_R·mae_R)/n
            // (`MAE.node_impurity`/`children_impurity`, L1 around each child
            // median, `_criterion.pyx:1450-1519`).
            RegressionCriterion::AbsoluteError => {
                let left_mae = mae_for_indices(data.y, &left_indices);
                let right_mae = mae_for_indices(data.y, &right_indices);
                let weighted_child_mae = (left_n_f * left_mae + right_n_f * right_mae) / n_f;
                parent_impurity - weighted_child_mae
            }
            // poisson: parent_deviance − (n_L·dev_L + n_R·dev_R)/n
            // (`Poisson.poisson_loss`, `_criterion.pyx:1671-1708`). A child with a
            // non-positive sum yields +∞ deviance ⇒ −∞ score ⇒ never selected,
            // mirroring sklearn's `y_sum <= EPSILON ⇒ INFINITY`.
            RegressionCriterion::Poisson => {
                let left_dev = poisson_deviance_for_indices(data.y, &left_indices);
                let right_dev = poisson_deviance_for_indices(data.y, &right_indices);
                let weighted_child_dev = (left_n_f * left_dev + right_n_f * right_dev) / n_f;
                parent_impurity - weighted_child_dev
            }
        };

        if score > best_score {
            best_score = score;
            best_feature = feat;
            best_threshold = threshold;
        }
    }

    if best_score > F::zero() {
        Some((best_feature, best_threshold, best_score * n_f))
    } else {
        None
    }
}

// ---------------------------------------------------------------------------
// Crate-internal functions for ensemble usage
// ---------------------------------------------------------------------------

/// Build a classification extra-tree with a subset of features for ensemble use.
///
/// Used internally by `ExtraTreesClassifier` to build individual trees.
#[allow(clippy::too_many_arguments)]
pub(crate) fn build_extra_classification_tree_for_ensemble<F: Float>(
    x: &Array2<F>,
    y: &[usize],
    n_classes: usize,
    indices: &[usize],
    feature_indices: Option<&[usize]>,
    params: &TreeParams,
    criterion: ClassificationCriterion,
    n_features: usize,
    max_features_n: usize,
    rng: &mut StdRng,
) -> Vec<Node<F>> {
    let data = ClassificationData {
        x,
        y,
        n_classes,
        feature_indices,
        criterion,
    };
    let mut nodes = Vec::new();
    build_extra_classification_tree(
        &data,
        indices,
        &mut nodes,
        0,
        params,
        n_features,
        max_features_n,
        rng,
    );
    nodes
}

/// Build a regression extra-tree with a subset of features for ensemble use.
///
/// Used internally by `ExtraTreesRegressor` to build individual trees.
#[allow(clippy::too_many_arguments)]
pub(crate) fn build_extra_regression_tree_for_ensemble<F: Float>(
    x: &Array2<F>,
    y: &Array1<F>,
    indices: &[usize],
    feature_indices: Option<&[usize]>,
    params: &TreeParams,
    n_features: usize,
    max_features_n: usize,
    rng: &mut StdRng,
) -> Vec<Node<F>> {
    // The ensemble path (`ExtraTreesRegressor`) currently builds MSE/mean trees
    // exclusively; preserve that exact behavior. Per-criterion ensemble support
    // is a separate unit (the ensemble builder signature lives across files).
    let data = RegressionData {
        x,
        y,
        feature_indices,
        criterion: RegressionCriterion::Mse,
    };
    let mut nodes = Vec::new();
    build_extra_regression_tree(
        &data,
        indices,
        &mut nodes,
        0,
        params,
        n_features,
        max_features_n,
        rng,
    );
    nodes
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_relative_eq;
    use ndarray::array;

    // -- Classifier tests --

    #[test]
    fn test_extra_classifier_simple_binary() {
        let x = Array2::from_shape_vec(
            (6, 2),
            vec![1.0, 2.0, 2.0, 3.0, 3.0, 3.0, 5.0, 6.0, 6.0, 7.0, 7.0, 8.0],
        )
        .unwrap();
        let y = array![0, 0, 0, 1, 1, 1];

        let model = ExtraTreeClassifier::<f64>::new().with_random_state(42);
        let fitted = model.fit(&x, &y).unwrap();
        let preds = fitted.predict(&x).unwrap();

        assert_eq!(preds.len(), 6);
        // ExtraTrees should separate linearly separable data.
        for i in 0..3 {
            assert_eq!(preds[i], 0, "sample {i} should be class 0");
        }
        for i in 3..6 {
            assert_eq!(preds[i], 1, "sample {i} should be class 1");
        }
    }

    #[test]
    fn test_extra_classifier_single_class() {
        let x = Array2::from_shape_vec((3, 1), vec![1.0, 2.0, 3.0]).unwrap();
        let y = array![0, 0, 0];

        let model = ExtraTreeClassifier::<f64>::new().with_random_state(42);
        let fitted = model.fit(&x, &y).unwrap();
        let preds = fitted.predict(&x).unwrap();

        assert_eq!(preds, array![0, 0, 0]);
    }

    #[test]
    fn test_extra_classifier_max_depth_1() {
        let x =
            Array2::from_shape_vec((8, 1), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]).unwrap();
        let y = array![0, 0, 0, 0, 1, 1, 1, 1];

        let model = ExtraTreeClassifier::<f64>::new()
            .with_max_depth(Some(1))
            .with_max_features(MaxFeatures::All)
            .with_random_state(42);
        let fitted = model.fit(&x, &y).unwrap();
        let _preds = fitted.predict(&x).unwrap();

        // With depth 1 and a single feature, it should still separate the classes.
        // The tree has exactly one split node and two leaves.
        assert_eq!(fitted.nodes().len(), 3);
    }

    #[test]
    fn test_extra_classifier_predict_proba() {
        let x = Array2::from_shape_vec((6, 1), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        let y = array![0, 0, 0, 1, 1, 1];

        let model = ExtraTreeClassifier::<f64>::new()
            .with_max_features(MaxFeatures::All)
            .with_random_state(42);
        let fitted = model.fit(&x, &y).unwrap();
        let proba = fitted.predict_proba(&x).unwrap();

        assert_eq!(proba.dim(), (6, 2));
        // Each row sums to 1.
        for i in 0..6 {
            let row_sum = proba.row(i).sum();
            assert_relative_eq!(row_sum, 1.0, epsilon = 1e-10);
        }
    }

    #[test]
    fn test_extra_classifier_feature_importances() {
        let x = Array2::from_shape_vec(
            (8, 2),
            vec![
                1.0, 1.0, 2.0, 1.0, 3.0, 1.0, 4.0, 1.0, 5.0, 1.0, 6.0, 1.0, 7.0, 1.0, 8.0, 1.0,
            ],
        )
        .unwrap();
        let y = array![0, 0, 0, 0, 1, 1, 1, 1];

        let model = ExtraTreeClassifier::<f64>::new()
            .with_max_features(MaxFeatures::All)
            .with_random_state(42);
        let fitted = model.fit(&x, &y).unwrap();
        let importances = fitted.feature_importances();

        assert_eq!(importances.len(), 2);
        // The sum of importances should be 1 (normalised).
        let total: f64 = importances.sum();
        assert_relative_eq!(total, 1.0, epsilon = 1e-10);
        // Feature 0 should have higher importance (feature 1 is constant).
        assert!(importances[0] > importances[1]);
    }

    #[test]
    fn test_extra_classifier_shape_mismatch() {
        let x = Array2::from_shape_vec((3, 2), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        let y = array![0, 0]; // wrong length

        let model = ExtraTreeClassifier::<f64>::new();
        assert!(model.fit(&x, &y).is_err());
    }

    #[test]
    fn test_extra_classifier_empty_data() {
        let x = Array2::<f64>::zeros((0, 2));
        let y = Array1::<usize>::zeros(0);

        let model = ExtraTreeClassifier::<f64>::new();
        assert!(model.fit(&x, &y).is_err());
    }

    #[test]
    fn test_extra_classifier_invalid_min_samples_split() {
        let x = Array2::from_shape_vec((3, 1), vec![1.0, 2.0, 3.0]).unwrap();
        let y = array![0, 0, 1];

        let model = ExtraTreeClassifier::<f64>::new().with_min_samples_split(1);
        assert!(model.fit(&x, &y).is_err());
    }

    #[test]
    fn test_extra_classifier_classes() {
        let x = Array2::from_shape_vec((6, 1), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        let y = array![0, 0, 0, 2, 2, 2]; // non-contiguous classes

        let model = ExtraTreeClassifier::<f64>::new().with_random_state(42);
        let fitted = model.fit(&x, &y).unwrap();

        assert_eq!(fitted.classes(), &[0, 2]);
        assert_eq!(fitted.n_classes(), 2);
    }

    #[test]
    fn test_extra_classifier_predict_shape_mismatch() {
        let x =
            Array2::from_shape_vec((4, 2), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]).unwrap();
        let y = array![0, 0, 1, 1];

        let model = ExtraTreeClassifier::<f64>::new().with_random_state(42);
        let fitted = model.fit(&x, &y).unwrap();

        let x_wrong = Array2::from_shape_vec((2, 3), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        assert!(fitted.predict(&x_wrong).is_err());
    }

    #[test]
    fn test_extra_classifier_f32() {
        let x = Array2::from_shape_vec(
            (6, 2),
            vec![
                1.0f32, 2.0, 2.0, 3.0, 3.0, 3.0, 5.0, 6.0, 6.0, 7.0, 7.0, 8.0,
            ],
        )
        .unwrap();
        let y = array![0, 0, 0, 1, 1, 1];

        let model = ExtraTreeClassifier::<f32>::new().with_random_state(42);
        let fitted = model.fit(&x, &y).unwrap();
        let preds = fitted.predict(&x).unwrap();
        assert_eq!(preds.len(), 6);
    }

    #[test]
    fn test_extra_classifier_deterministic() {
        let x = Array2::from_shape_vec(
            (8, 2),
            vec![
                1.0, 2.0, 2.0, 3.0, 3.0, 3.0, 4.0, 4.0, 5.0, 6.0, 6.0, 7.0, 7.0, 8.0, 8.0, 9.0,
            ],
        )
        .unwrap();
        let y = array![0, 0, 0, 0, 1, 1, 1, 1];

        let model1 = ExtraTreeClassifier::<f64>::new().with_random_state(123);
        let model2 = ExtraTreeClassifier::<f64>::new().with_random_state(123);

        let fitted1 = model1.fit(&x, &y).unwrap();
        let fitted2 = model2.fit(&x, &y).unwrap();

        let preds1 = fitted1.predict(&x).unwrap();
        let preds2 = fitted2.predict(&x).unwrap();

        assert_eq!(preds1, preds2);
    }

    // -- Regressor tests --

    #[test]
    fn test_extra_regressor_simple() {
        let x = Array2::from_shape_vec((6, 1), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        let y = array![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];

        let model = ExtraTreeRegressor::<f64>::new()
            .with_max_features(MaxFeatures::All)
            .with_random_state(42);
        let fitted = model.fit(&x, &y).unwrap();
        let preds = fitted.predict(&x).unwrap();

        // A deep extra-tree should roughly memorize the training data.
        assert_eq!(preds.len(), 6);
        for i in 0..6 {
            assert_relative_eq!(preds[i], y[i], epsilon = 1.0);
        }
    }

    #[test]
    fn test_extra_regressor_constant_target() {
        let x = Array2::from_shape_vec((4, 1), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
        let y = array![5.0, 5.0, 5.0, 5.0];

        let model = ExtraTreeRegressor::<f64>::new().with_random_state(42);
        let fitted = model.fit(&x, &y).unwrap();
        let preds = fitted.predict(&x).unwrap();

        for &p in &preds {
            assert_relative_eq!(p, 5.0, epsilon = 1e-10);
        }
    }

    #[test]
    fn test_extra_regressor_feature_importances() {
        let x = Array2::from_shape_vec(
            (8, 2),
            vec![
                1.0, 0.0, 2.0, 0.0, 3.0, 0.0, 4.0, 0.0, 5.0, 0.0, 6.0, 0.0, 7.0, 0.0, 8.0, 0.0,
            ],
        )
        .unwrap();
        let y = array![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];

        let model = ExtraTreeRegressor::<f64>::new()
            .with_max_features(MaxFeatures::All)
            .with_random_state(42);
        let fitted = model.fit(&x, &y).unwrap();
        let importances = fitted.feature_importances();

        assert_eq!(importances.len(), 2);
        let total: f64 = importances.sum();
        assert_relative_eq!(total, 1.0, epsilon = 1e-10);
        // Feature 0 drives the target; feature 1 is constant.
        assert!(importances[0] > importances[1]);
    }

    #[test]
    fn test_extra_regressor_shape_mismatch() {
        let x = Array2::from_shape_vec((3, 2), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        let y = array![1.0, 2.0]; // wrong length

        let model = ExtraTreeRegressor::<f64>::new();
        assert!(model.fit(&x, &y).is_err());
    }

    #[test]
    fn test_extra_regressor_empty_data() {
        let x = Array2::<f64>::zeros((0, 2));
        let y = Array1::<f64>::zeros(0);

        let model = ExtraTreeRegressor::<f64>::new();
        assert!(model.fit(&x, &y).is_err());
    }

    #[test]
    fn test_extra_regressor_predict_shape_mismatch() {
        let x =
            Array2::from_shape_vec((4, 2), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]).unwrap();
        let y = array![1.0, 2.0, 3.0, 4.0];

        let model = ExtraTreeRegressor::<f64>::new().with_random_state(42);
        let fitted = model.fit(&x, &y).unwrap();

        let x_wrong = Array2::from_shape_vec((2, 3), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        assert!(fitted.predict(&x_wrong).is_err());
    }

    #[test]
    fn test_extra_regressor_max_depth() {
        let x =
            Array2::from_shape_vec((8, 1), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]).unwrap();
        let y = array![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];

        let model = ExtraTreeRegressor::<f64>::new()
            .with_max_depth(Some(1))
            .with_max_features(MaxFeatures::All)
            .with_random_state(42);
        let fitted = model.fit(&x, &y).unwrap();

        // With depth 1, the tree should have exactly 3 nodes: one split + two leaves.
        assert_eq!(fitted.nodes().len(), 3);
    }

    #[test]
    fn test_extra_regressor_deterministic() {
        let x = Array2::from_shape_vec((6, 1), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        let y = array![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];

        let model1 = ExtraTreeRegressor::<f64>::new().with_random_state(99);
        let model2 = ExtraTreeRegressor::<f64>::new().with_random_state(99);

        let fitted1 = model1.fit(&x, &y).unwrap();
        let fitted2 = model2.fit(&x, &y).unwrap();

        let preds1 = fitted1.predict(&x).unwrap();
        let preds2 = fitted2.predict(&x).unwrap();

        for i in 0..6 {
            assert_relative_eq!(preds1[i], preds2[i], epsilon = 1e-12);
        }
    }

    #[test]
    fn test_extra_regressor_f32() {
        let x = Array2::from_shape_vec((4, 1), vec![1.0f32, 2.0, 3.0, 4.0]).unwrap();
        let y = array![1.0f32, 2.0, 3.0, 4.0];

        let model = ExtraTreeRegressor::<f32>::new().with_random_state(42);
        let fitted = model.fit(&x, &y).unwrap();
        let preds = fitted.predict(&x).unwrap();
        assert_eq!(preds.len(), 4);
    }

    // -- Builder tests --

    #[test]
    fn test_classifier_builder_methods() {
        let model = ExtraTreeClassifier::<f64>::new()
            .with_max_depth(Some(5))
            .with_min_samples_split(10)
            .with_min_samples_leaf(3)
            .with_max_features(MaxFeatures::Log2)
            .with_criterion(ClassificationCriterion::Entropy)
            .with_random_state(42);

        assert_eq!(model.max_depth, Some(5));
        assert_eq!(model.min_samples_split, 10);
        assert_eq!(model.min_samples_leaf, 3);
        assert_eq!(model.max_features, MaxFeatures::Log2);
        assert_eq!(model.criterion, ClassificationCriterion::Entropy);
        assert_eq!(model.random_state, Some(42));
    }

    #[test]
    fn test_regressor_builder_methods() {
        let model = ExtraTreeRegressor::<f64>::new()
            .with_max_depth(Some(10))
            .with_min_samples_split(5)
            .with_min_samples_leaf(2)
            .with_max_features(MaxFeatures::Fixed(3))
            .with_criterion(RegressionCriterion::Mse)
            .with_random_state(99);

        assert_eq!(model.max_depth, Some(10));
        assert_eq!(model.min_samples_split, 5);
        assert_eq!(model.min_samples_leaf, 2);
        assert_eq!(model.max_features, MaxFeatures::Fixed(3));
        assert_eq!(model.criterion, RegressionCriterion::Mse);
        assert_eq!(model.random_state, Some(99));
    }

    #[test]
    fn test_classifier_default() {
        let model = ExtraTreeClassifier::<f64>::default();
        assert_eq!(model.max_depth, None);
        assert_eq!(model.min_samples_split, 2);
        assert_eq!(model.min_samples_leaf, 1);
        assert_eq!(model.max_features, MaxFeatures::Sqrt);
        assert_eq!(model.criterion, ClassificationCriterion::Gini);
        assert_eq!(model.random_state, None);
    }

    #[test]
    fn test_regressor_default() {
        let model = ExtraTreeRegressor::<f64>::default();
        assert_eq!(model.max_depth, None);
        assert_eq!(model.min_samples_split, 2);
        assert_eq!(model.min_samples_leaf, 1);
        assert_eq!(model.max_features, MaxFeatures::All);
        assert_eq!(model.criterion, RegressionCriterion::Mse);
        assert_eq!(model.random_state, None);
    }
}