perpetual 3.0.0-rc.2

//! Prediction Methods
//!
//! Prediction, probability, contribution, and feature importance methods for the booster.

use crate::MultiOutputBooster;

use crate::booster::config::{CalibrationMethod, ContributionsMethod};
use crate::data::ColumnarMatrix;
use crate::objective::Objective;
use crate::{Matrix, PerpetualBooster, shapley::predict_contributions_row_shapley, tree::core::Tree, utils::odds};
use rayon::prelude::*;
use std::collections::{HashMap, HashSet};

impl PerpetualBooster {
    pub(crate) fn predict_tree_ensemble(&self, data: &Matrix<f64>, parallel: bool) -> Vec<f64> {
        let mut init_preds = vec![self.base_score; data.rows];
        self.get_prediction_trees().iter().for_each(|tree| {
            for (prediction, value) in init_preds
                .iter_mut()
                .zip(tree.predict(data, parallel, &self.cfg.missing))
            {
                *prediction += value;
            }
        });
        init_preds
    }

    pub(crate) fn predict_tree_ensemble_columnar(&self, data: &ColumnarMatrix<f64>, parallel: bool) -> Vec<f64> {
        let mut init_preds = vec![self.base_score; data.rows];
        self.get_prediction_trees().iter().for_each(|tree| {
            for (prediction, value) in
                init_preds
                    .iter_mut()
                    .zip(tree.predict_columnar(data, parallel, &self.cfg.missing))
            {
                *prediction += value;
            }
        });
        init_preds
    }

    fn apply_regression_head(&self, data: &Matrix<f64>, predictions: &mut [f64]) {
        let Some(head) = self.regression_head.as_ref() else {
            return;
        };

        for (row_idx, prediction) in predictions.iter_mut().enumerate() {
            let mut residual_prediction = head.intercept;
            for col_idx in 0..data.cols {
                let value = *data.get(row_idx, col_idx);
                let standardized = if value.is_finite() {
                    (value - head.feature_means[col_idx]) / head.feature_scales[col_idx]
                } else {
                    0.0
                };
                residual_prediction += head.coefficients[col_idx] * standardized;
            }
            *prediction += head.blend_weight * residual_prediction;
        }
    }

    fn apply_regression_head_columnar(&self, data: &ColumnarMatrix<f64>, predictions: &mut [f64]) {
        let Some(head) = self.regression_head.as_ref() else {
            return;
        };

        for (row_idx, prediction) in predictions.iter_mut().enumerate() {
            let mut residual_prediction = head.intercept;
            for col_idx in 0..data.cols {
                let standardized = if data.is_valid(row_idx, col_idx) {
                    let value = *data.get(row_idx, col_idx);
                    if value.is_finite() {
                        (value - head.feature_means[col_idx]) / head.feature_scales[col_idx]
                    } else {
                        0.0
                    }
                } else {
                    0.0
                };
                residual_prediction += head.coefficients[col_idx] * standardized;
            }
            *prediction += head.blend_weight * residual_prediction;
        }
    }

    fn predict_fold_logits(&self, data: &Matrix<f64>, parallel: bool) -> Vec<[f64; 5]> {
        let mut results = vec![[self.base_score; 5]; data.rows];

        for tree in self.get_prediction_trees() {
            let tree_weights = tree.predict_weights(data, parallel, &self.cfg.missing);
            for (row_logits, row_weights) in results.iter_mut().zip(tree_weights.iter()) {
                for fold in 0..5 {
                    row_logits[fold] += row_weights[fold] as f64;
                }
            }
        }

        results
    }

    fn predict_fold_logits_columnar(&self, data: &ColumnarMatrix<f64>, parallel: bool) -> Vec<[f64; 5]> {
        let mut results = vec![[self.base_score; 5]; data.rows];

        for tree in self.get_prediction_trees() {
            let tree_weights = tree.predict_weights_columnar(data, parallel, &self.cfg.missing);
            for (row_logits, row_weights) in results.iter_mut().zip(tree_weights.iter()) {
                for fold in 0..5 {
                    row_logits[fold] += row_weights[fold] as f64;
                }
            }
        }

        results
    }

    /// Generate predictions for the given data.
    ///
    /// # Arguments
    ///
    /// * `data` - The feature matrix.
    /// * `parallel` - If `true`, predictions are computed in parallel using Rayon.
    ///
    /// # Returns
    ///
    /// A vector of predicted values.
    /// * For regression/ranking: The raw predicted score.
    /// * For classification: The log-odds (logit). Apply sigmoid to get probabilities.
    pub fn predict(&self, data: &Matrix<f64>, parallel: bool) -> Vec<f64> {
        let mut predictions = self.predict_tree_ensemble(data, parallel);
        self.apply_regression_head(data, &mut predictions);
        predictions
    }

    /// Generate predictions for columnar data (Zero-Copy).
    ///
    /// # Arguments
    ///
    /// * `data` - The columnar feature matrix.
    /// * `parallel` - If `true`, predictions are computed in parallel.
    pub fn predict_columnar(&self, data: &ColumnarMatrix<f64>, parallel: bool) -> Vec<f64> {
        let mut predictions = self.predict_tree_ensemble_columnar(data, parallel);
        self.apply_regression_head_columnar(data, &mut predictions);
        predictions
    }

    /// Generate class probabilities (sigmoid of log-odds) for the given data.
    ///
    /// Only meaningful for binary classification with `LogLoss` objective.
    ///
    /// * `data` - The feature matrix.
    /// * `parallel` - Predict in parallel.
    pub fn predict_proba(&self, data: &Matrix<f64>, parallel: bool, calibrated: bool) -> Vec<f64> {
        let mut proba: Vec<f64> = if matches!(self.cfg.objective, Objective::LogLoss) {
            let fold_logits = self.predict_fold_logits(data, parallel);
            if parallel {
                fold_logits
                    .par_iter()
                    .map(|row| row.iter().map(|p| odds(*p)).sum::<f64>() / 5.0)
                    .collect()
            } else {
                fold_logits
                    .iter()
                    .map(|row| row.iter().map(|p| odds(*p)).sum::<f64>() / 5.0)
                    .collect()
            }
        } else {
            let preds = self.predict(data, parallel);
            if parallel {
                preds.par_iter().map(|p| odds(*p)).collect()
            } else {
                preds.iter().map(|p| odds(*p)).collect()
            }
        };

        if let Some(calibrator) = self.isotonic_calibrator.as_ref().filter(|_| calibrated) {
            let scores = match self.cfg.calibration_method {
                CalibrationMethod::Conformal => proba.clone(),
                _ => self.get_calibration_scores(data, parallel),
            };
            proba = calibrator.transform(&scores);
        }
        proba
    }

    /// Generate class probabilities for columnar data (zero-copy).
    ///
    /// * `data` - The columnar feature matrix.
    /// * `parallel` - Predict in parallel.
    pub fn predict_proba_columnar(&self, data: &ColumnarMatrix<f64>, parallel: bool, calibrated: bool) -> Vec<f64> {
        let mut proba: Vec<f64> = if matches!(self.cfg.objective, Objective::LogLoss) {
            let fold_logits = self.predict_fold_logits_columnar(data, parallel);
            if parallel {
                fold_logits
                    .par_iter()
                    .map(|row| row.iter().map(|p| odds(*p)).sum::<f64>() / 5.0)
                    .collect()
            } else {
                fold_logits
                    .iter()
                    .map(|row| row.iter().map(|p| odds(*p)).sum::<f64>() / 5.0)
                    .collect()
            }
        } else {
            let preds = self.predict_columnar(data, parallel);
            if parallel {
                preds.par_iter().map(|p| odds(*p)).collect()
            } else {
                preds.iter().map(|p| odds(*p)).collect()
            }
        };

        if let Some(calibrator) = self.isotonic_calibrator.as_ref().filter(|_| calibrated) {
            let scores = match self.cfg.calibration_method {
                CalibrationMethod::Conformal => proba.clone(),
                _ => self.get_calibration_scores_columnar(data, parallel),
            };
            proba = calibrator.transform(&scores);
        }
        proba
    }

    /// Calculate Feature Contributions (SHAP-like values).
    ///
    /// Computes how much each feature contributed to the prediction for each sample.
    /// - The sum of contributions + bias equals the prediction.
    /// - The last column in the returned matrix is the bias term.
    ///
    /// # Arguments
    ///
    /// * `data` - Feature matrix.
    /// * `method` - Method to calculate contributions (`Average` is the standard TreeSHAP approximation).
    /// * `parallel` - Run in parallel.
    pub fn predict_contributions(&self, data: &Matrix<f64>, method: ContributionsMethod, parallel: bool) -> Vec<f64> {
        match method {
            ContributionsMethod::Average => self.predict_contributions_average(data, parallel),
            ContributionsMethod::ProbabilityChange => {
                match self.cfg.objective {
                    Objective::LogLoss => {}
                    _ => panic!("ProbabilityChange contributions method is only valid when LogLoss objective is used."),
                }
                self.predict_contributions_probability_change(data, parallel)
            }
            _ => self.predict_contributions_tree_alone(data, parallel, method),
        }
    }

    // All of the contribution calculation methods, except for average are calculated
    // using just the model, so we don't need to have separate methods, we can instead
    // just have this one method, that dispatches to each one respectively.
    fn predict_contributions_tree_alone(
        &self,
        data: &Matrix<f64>,
        parallel: bool,
        method: ContributionsMethod,
    ) -> Vec<f64> {
        let mut contribs = vec![0.; (data.cols + 1) * data.rows];

        // Add the bias term to every bias value...
        let bias_idx = data.cols + 1;
        contribs
            .iter_mut()
            .skip(bias_idx - 1)
            .step_by(bias_idx)
            .for_each(|v| *v += self.base_score);

        let row_pred_fn = match method {
            ContributionsMethod::Weight => Tree::predict_contributions_row_weight,
            ContributionsMethod::BranchDifference => Tree::predict_contributions_row_branch_difference,
            ContributionsMethod::MidpointDifference => Tree::predict_contributions_row_midpoint_difference,
            ContributionsMethod::ModeDifference => Tree::predict_contributions_row_mode_difference,
            ContributionsMethod::Shapley => predict_contributions_row_shapley,
            ContributionsMethod::Average | ContributionsMethod::ProbabilityChange => unreachable!(),
        };
        // Clean this up..
        // materializing a row, and then passing that to all of the
        // trees seems to be the fastest approach (5X faster), we should test
        // something like this for normal predictions.
        if parallel {
            data.index
                .par_iter()
                .zip(contribs.par_chunks_mut(data.cols + 1))
                .for_each(|(row, c)| {
                    let r_ = data.get_row(*row);
                    self.get_prediction_trees().iter().for_each(|t| {
                        row_pred_fn(t, &r_, c, &self.cfg.missing);
                    });
                });
        } else {
            data.index
                .iter()
                .zip(contribs.chunks_mut(data.cols + 1))
                .for_each(|(row, c)| {
                    let r_ = data.get_row(*row);
                    self.get_prediction_trees().iter().for_each(|t| {
                        row_pred_fn(t, &r_, c, &self.cfg.missing);
                    });
                });
        }

        contribs
    }

    /// Generate predictions on data using the gradient booster.
    ///
    /// * `data` -  Either a Polars or Pandas DataFrame, or a 2 dimensional Numpy array.
    fn predict_contributions_average(&self, data: &Matrix<f64>, parallel: bool) -> Vec<f64> {
        let weights: Vec<HashMap<usize, f64>> = if parallel {
            self.get_prediction_trees()
                .par_iter()
                .map(|t| t.distribute_leaf_weights())
                .collect()
        } else {
            self.get_prediction_trees()
                .iter()
                .map(|t| t.distribute_leaf_weights())
                .collect()
        };
        let mut contribs = vec![0.0; (data.cols + 1) * data.rows];

        // Add the bias term to every bias value...
        let bias_idx = data.cols + 1;
        contribs
            .iter_mut()
            .skip(bias_idx - 1)
            .step_by(bias_idx)
            .for_each(|v| *v += self.base_score);

        // Clean this up..
        // materializing a row, and then passing that to all of the
        // trees seems to be the fastest approach (5X faster), we should test
        // something like this for normal predictions.
        if parallel {
            data.index
                .par_iter()
                .zip(contribs.par_chunks_mut(data.cols + 1))
                .for_each(|(row, c)| {
                    let r_ = data.get_row(*row);
                    self.get_prediction_trees()
                        .iter()
                        .zip(weights.iter())
                        .for_each(|(t, w)| {
                            t.predict_contributions_row_average(&r_, c, w, &self.cfg.missing);
                        });
                });
        } else {
            data.index
                .iter()
                .zip(contribs.chunks_mut(data.cols + 1))
                .for_each(|(row, c)| {
                    let r_ = data.get_row(*row);
                    self.get_prediction_trees()
                        .iter()
                        .zip(weights.iter())
                        .for_each(|(t, w)| {
                            t.predict_contributions_row_average(&r_, c, w, &self.cfg.missing);
                        });
                });
        }

        contribs
    }

    fn predict_contributions_probability_change(&self, data: &Matrix<f64>, parallel: bool) -> Vec<f64> {
        let mut contribs = vec![0.; (data.cols + 1) * data.rows];
        let bias_idx = data.cols + 1;
        contribs
            .iter_mut()
            .skip(bias_idx - 1)
            .step_by(bias_idx)
            .for_each(|v| *v += odds(self.base_score));

        if parallel {
            data.index
                .par_iter()
                .zip(contribs.par_chunks_mut(data.cols + 1))
                .for_each(|(row, c)| {
                    let r_ = data.get_row(*row);
                    self.get_prediction_trees().iter().fold(self.base_score, |acc, t| {
                        t.predict_contributions_row_probability_change(&r_, c, &self.cfg.missing, acc)
                    });
                });
        } else {
            data.index
                .iter()
                .zip(contribs.chunks_mut(data.cols + 1))
                .for_each(|(row, c)| {
                    let r_ = data.get_row(*row);
                    self.get_prediction_trees().iter().fold(self.base_score, |acc, t| {
                        t.predict_contributions_row_probability_change(&r_, c, &self.cfg.missing, acc)
                    });
                });
        }
        contribs
    }

    /// Generate node predictions on data using the gradient booster.
    ///
    /// * `data` -  Either a Polars or Pandas DataFrame, or a 2 dimensional Numpy array.
    /// * `parallel` -  Predict in parallel.
    pub fn predict_nodes(&self, data: &Matrix<f64>, parallel: bool) -> Vec<Vec<HashSet<usize>>> {
        let mut v = Vec::with_capacity(data.rows);
        self.get_prediction_trees().iter().for_each(|tree| {
            let tree_nodes = tree.predict_nodes(data, parallel, &self.cfg.missing);
            v.push(tree_nodes);
        });
        v
    }

    /// Generate node predictions on columnar data using the gradient booster (zero-copy from Polars).
    pub fn predict_nodes_columnar(&self, data: &ColumnarMatrix<f64>, parallel: bool) -> Vec<Vec<HashSet<usize>>> {
        let mut v = Vec::with_capacity(data.rows);
        self.get_prediction_trees().iter().for_each(|tree| {
            let tree_nodes = tree.predict_nodes_columnar(data, parallel, &self.cfg.missing);
            v.push(tree_nodes);
        });
        v
    }

    fn predict_contributions_average_columnar(&self, data: &ColumnarMatrix<f64>, parallel: bool) -> Vec<f64> {
        let weights: Vec<HashMap<usize, f64>> = if parallel {
            self.get_prediction_trees()
                .par_iter()
                .map(|t| t.distribute_leaf_weights())
                .collect()
        } else {
            self.get_prediction_trees()
                .iter()
                .map(|t| t.distribute_leaf_weights())
                .collect()
        };
        let mut contribs = vec![0.; (data.cols + 1) * data.rows];
        let bias_idx = data.cols + 1;
        contribs
            .iter_mut()
            .skip(bias_idx - 1)
            .step_by(bias_idx)
            .for_each(|v| *v += self.base_score);

        if parallel {
            data.index
                .par_iter()
                .zip(contribs.par_chunks_mut(data.cols + 1))
                .for_each(|(row, c)| {
                    let r_ = data.get_row(*row);
                    self.get_prediction_trees()
                        .iter()
                        .zip(weights.iter())
                        .for_each(|(t, w)| {
                            t.predict_contributions_row_average(&r_, c, w, &self.cfg.missing);
                        });
                });
        } else {
            data.index
                .iter()
                .zip(contribs.chunks_mut(data.cols + 1))
                .for_each(|(row, c)| {
                    let r_ = data.get_row(*row);
                    self.get_prediction_trees()
                        .iter()
                        .zip(weights.iter())
                        .for_each(|(t, w)| {
                            t.predict_contributions_row_average(&r_, c, w, &self.cfg.missing);
                        });
                });
        }

        contribs
    }

    fn predict_contributions_probability_change_columnar(
        &self,
        data: &ColumnarMatrix<f64>,
        parallel: bool,
    ) -> Vec<f64> {
        let mut contribs = vec![0.; (data.cols + 1) * data.rows];
        let bias_idx = data.cols + 1;
        contribs
            .iter_mut()
            .skip(bias_idx - 1)
            .step_by(bias_idx)
            .for_each(|v| *v += odds(self.base_score));

        if parallel {
            data.index
                .par_iter()
                .zip(contribs.par_chunks_mut(data.cols + 1))
                .for_each(|(row, c)| {
                    let r_ = data.get_row(*row);
                    self.get_prediction_trees().iter().fold(self.base_score, |acc, t| {
                        t.predict_contributions_row_probability_change(&r_, c, &self.cfg.missing, acc)
                    });
                });
        } else {
            data.index
                .iter()
                .zip(contribs.chunks_mut(data.cols + 1))
                .for_each(|(row, c)| {
                    let r_ = data.get_row(*row);
                    self.get_prediction_trees().iter().fold(self.base_score, |acc, t| {
                        t.predict_contributions_row_probability_change(&r_, c, &self.cfg.missing, acc)
                    });
                });
        }
        contribs
    }

    fn predict_contributions_tree_alone_columnar(
        &self,
        data: &ColumnarMatrix<f64>,
        parallel: bool,
        method: ContributionsMethod,
    ) -> Vec<f64> {
        let mut contribs = vec![0.; (data.cols + 1) * data.rows];
        let row_pred_fn = match method {
            ContributionsMethod::Weight => Tree::predict_contributions_row_weight,
            ContributionsMethod::BranchDifference => Tree::predict_contributions_row_branch_difference,
            ContributionsMethod::MidpointDifference => Tree::predict_contributions_row_midpoint_difference,
            ContributionsMethod::ModeDifference => Tree::predict_contributions_row_mode_difference,
            ContributionsMethod::Shapley => predict_contributions_row_shapley,
            ContributionsMethod::Average | ContributionsMethod::ProbabilityChange => unreachable!(),
        };
        if parallel {
            data.index
                .par_iter()
                .zip(contribs.par_chunks_mut(data.cols + 1))
                .for_each(|(row, c)| {
                    let r_ = data.get_row(*row);
                    self.get_prediction_trees().iter().for_each(|t| {
                        row_pred_fn(t, &r_, c, &self.cfg.missing);
                    });
                });
        } else {
            data.index
                .iter()
                .zip(contribs.chunks_mut(data.cols + 1))
                .for_each(|(row, c)| {
                    let r_ = data.get_row(*row);
                    self.get_prediction_trees().iter().for_each(|t| {
                        row_pred_fn(t, &r_, c, &self.cfg.missing);
                    });
                });
        }
        contribs
    }

    pub fn predict_contributions_columnar(
        &self,
        data: &ColumnarMatrix<f64>,
        method: ContributionsMethod,
        parallel: bool,
    ) -> Vec<f64> {
        match method {
            ContributionsMethod::Average => self.predict_contributions_average_columnar(data, parallel),
            ContributionsMethod::ProbabilityChange => {
                match self.cfg.objective {
                    Objective::LogLoss => {}
                    _ => panic!("ProbabilityChange contributions method is only valid when LogLoss objective is used."),
                }
                self.predict_contributions_probability_change_columnar(data, parallel)
            }
            _ => self.predict_contributions_tree_alone_columnar(data, parallel, method),
        }
    }

    /// Predict a distribution using uncalibrated leaf weights from internal nodes.
    /// Returns `n` predictions for each sample, resulting in a matrix of size `(n_samples, n)`.
    pub fn predict_distribution(&self, data: &Matrix<f64>, n: usize, parallel: bool) -> Vec<Vec<f64>> {
        use rand::RngExt;
        use rand::SeedableRng;
        let n_samples = data.rows;
        let mut results = vec![vec![self.base_score; n]; n_samples];

        for tree in &self.trees {
            let tree_weights = tree.predict_weights(data, parallel, &self.cfg.missing);
            for (i, row_weights) in tree_weights.iter().enumerate() {
                // Ensure reproducible results for each sample across different runs
                // by hashing the provided seed with the sample index and tree index.
                let mut rng = rand::rngs::StdRng::seed_from_u64(
                    self.cfg
                        .seed
                        .wrapping_add(i as u64)
                        .wrapping_add(tree.nodes.len() as u64),
                );
                for val in results[i].iter_mut().take(n) {
                    let idx = rng.random_range(0..5);
                    *val += row_weights[idx] as f64;
                }
            }
        }
        // Sanitize any non-finite values (defensive: replace NaN/inf with base_score)
        for row in results.iter_mut() {
            for v in row.iter_mut() {
                if !v.is_finite() {
                    *v = self.base_score;
                }
            }
        }
        results
    }

    pub fn predict_distribution_columnar(&self, data: &ColumnarMatrix<f64>, n: usize, parallel: bool) -> Vec<Vec<f64>> {
        use rand::RngExt;
        use rand::SeedableRng;
        let n_samples = data.index.len();
        let mut results = vec![vec![self.base_score; n]; n_samples];

        for tree in &self.trees {
            let tree_weights = tree.predict_weights_columnar(data, parallel, &self.cfg.missing);
            for (i, row_weights) in tree_weights.iter().enumerate() {
                let mut rng = rand::rngs::StdRng::seed_from_u64(
                    self.cfg
                        .seed
                        .wrapping_add(i as u64)
                        .wrapping_add(tree.nodes.len() as u64),
                );
                for val in results[i].iter_mut().take(n) {
                    let idx = rng.random_range(0..5);
                    *val += row_weights[idx] as f64;
                }
            }
        }
        // Sanitize any non-finite values (defensive: replace NaN/inf with base_score)
        for row in results.iter_mut() {
            for v in row.iter_mut() {
                if !v.is_finite() {
                    *v = self.base_score;
                }
            }
        }
        results
    }
}

impl MultiOutputBooster {
    fn softmax_probabilities_with_temperature(logits: &[f64], temperature: f64) -> Vec<f64> {
        let max_logit = logits.iter().copied().fold(f64::NEG_INFINITY, f64::max);
        let scale = temperature.max(f64::EPSILON);
        let exp_logits = logits
            .iter()
            .map(|value| ((value - max_logit) / scale).exp())
            .collect::<Vec<_>>();
        let normalizer = exp_logits.iter().sum::<f64>().max(f64::EPSILON);
        exp_logits.into_iter().map(|value| value / normalizer).collect()
    }

    fn native_multiclass_temperature(&self) -> f64 {
        if !self.native_multiclass {
            return 1.0;
        }

        match self.n_boosters {
            0..=4 => 1.0,
            5 => 1.5,
            6..=8 => 2.0,
            _ => 1.75,
        }
    }

    fn multiclass_probability_alpha(&self) -> f64 {
        if self.n_boosters == 3 {
            0.25
        } else if (4..=5).contains(&self.n_boosters) {
            0.5
        } else {
            0.0
        }
    }

    fn multiclass_probability_beta(&self) -> f64 {
        0.25
    }

    fn normalize_multiclass_probabilities(probabilities: &mut [f64], alpha: f64, priors: Option<&[f64]>, beta: f64) {
        probabilities.iter_mut().enumerate().for_each(|(idx, value)| {
            let clipped = value.clamp(f64::EPSILON, 1.0 - f64::EPSILON);
            let prior = priors
                .and_then(|class_priors| class_priors.get(idx))
                .copied()
                .unwrap_or(1.0)
                .clamp(f64::EPSILON, 1.0);
            *value = clipped / (1.0 - clipped).powf(alpha) / prior.powf(beta);
        });

        let sum = probabilities.iter().sum::<f64>();
        if sum <= f64::EPSILON {
            let uniform = 1.0 / probabilities.len().max(1) as f64;
            probabilities.iter_mut().for_each(|value| *value = uniform);
            return;
        }

        probabilities.iter_mut().for_each(|value| *value /= sum);
    }

    /// Generate predictions on data using the multi-output booster.
    ///
    /// * `data` -  Either a Polars or Pandas DataFrame, or a 2 dimensional Numpy array.
    /// * `parallel` -  Predict in parallel.
    pub fn predict(&self, data: &Matrix<f64>, parallel: bool) -> Vec<f64> {
        self.boosters
            .iter()
            .flat_map(|b| b.predict(data, parallel))
            .collect::<Vec<f64>>()
    }

    /// Generate predictions on columnar data using the multi-output booster (zero-copy from Polars).
    pub fn predict_columnar(&self, data: &ColumnarMatrix<f64>, parallel: bool) -> Vec<f64> {
        self.boosters
            .iter()
            .flat_map(|b| b.predict_columnar(data, parallel))
            .collect::<Vec<f64>>()
    }

    /// Generate probabilities on data using the multi-output booster.
    ///
    /// * `data` -  Either a Polars or Pandas DataFrame, or a 2 dimensional Numpy array.
    /// * `parallel` -  Predict in parallel.
    pub fn predict_proba(&self, data: &Matrix<f64>, parallel: bool) -> Vec<f64> {
        if self.native_multiclass {
            let logits = self.predict(data, parallel);
            let logits_matrix = Matrix::new(&logits, data.rows, self.n_boosters);
            let mut preds = Vec::with_capacity(logits.len());
            let temperature = self.native_multiclass_temperature();
            for row in 0..data.rows {
                let row_logits = logits_matrix.get_row(row);
                preds.extend(Self::softmax_probabilities_with_temperature(&row_logits, temperature));
            }
            return preds;
        }

        let alpha = self.multiclass_probability_alpha();
        let beta = self.multiclass_probability_beta();
        let class_probabilities = self
            .boosters
            .iter()
            .map(|b| b.predict_proba(data, parallel, false))
            .collect::<Vec<_>>();
        let mut preds = Vec::with_capacity(data.rows * self.n_boosters);
        for row in 0..data.rows {
            let mut probabilities = class_probabilities
                .iter()
                .map(|class_probs| class_probs[row])
                .collect::<Vec<f64>>();
            Self::normalize_multiclass_probabilities(&mut probabilities, alpha, Some(&self.class_priors), beta);
            preds.extend(probabilities);
        }
        preds
    }

    /// Generate probabilities on columnar data using the multi-output booster (zero-copy from Polars).
    pub fn predict_proba_columnar(&self, data: &ColumnarMatrix<f64>, parallel: bool) -> Vec<f64> {
        if self.native_multiclass {
            let logits = self.predict_columnar(data, parallel);
            let logits_matrix = Matrix::new(&logits, data.rows, self.n_boosters);
            let mut preds = Vec::with_capacity(logits.len());
            let temperature = self.native_multiclass_temperature();
            for row in 0..data.rows {
                let row_logits = logits_matrix.get_row(row);
                preds.extend(Self::softmax_probabilities_with_temperature(&row_logits, temperature));
            }
            return preds;
        }

        let alpha = self.multiclass_probability_alpha();
        let beta = self.multiclass_probability_beta();
        let class_probabilities = self
            .boosters
            .iter()
            .map(|b| b.predict_proba_columnar(data, parallel, false))
            .collect::<Vec<_>>();
        let mut preds = Vec::with_capacity(data.rows * self.n_boosters);
        for row in 0..data.rows {
            let mut probabilities = class_probabilities
                .iter()
                .map(|class_probs| class_probs[row])
                .collect::<Vec<f64>>();
            Self::normalize_multiclass_probabilities(&mut probabilities, alpha, Some(&self.class_priors), beta);
            preds.extend(probabilities);
        }
        preds
    }

    /// Generate node predictions on data using the gradient booster.
    ///
    /// * `data` -  Either a Polars or Pandas DataFrame, or a 2 dimensional Numpy array.
    /// * `parallel` -  Predict in parallel.
    pub fn predict_nodes(&self, data: &Matrix<f64>, parallel: bool) -> Vec<Vec<Vec<HashSet<usize>>>> {
        self.boosters.iter().map(|b| b.predict_nodes(data, parallel)).collect()
    }

    /// Generate node predictions on columnar data using the gradient booster (zero-copy from Polars).
    pub fn predict_nodes_columnar(&self, data: &ColumnarMatrix<f64>, parallel: bool) -> Vec<Vec<Vec<HashSet<usize>>>> {
        self.boosters
            .iter()
            .map(|b| b.predict_nodes_columnar(data, parallel))
            .collect()
    }

    /// Predict with the fitted booster on new data, returning the feature
    /// contribution matrix. The last column is the bias term.
    pub fn predict_contributions(&self, data: &Matrix<f64>, method: ContributionsMethod, parallel: bool) -> Vec<f64> {
        self.boosters
            .iter()
            .flat_map(|b| b.predict_contributions(data, method, parallel))
            .collect::<Vec<f64>>()
    }

    /// Calculate Feature Contributions for columnar data.
    pub fn predict_contributions_columnar(
        &self,
        data: &ColumnarMatrix<f64>,
        method: ContributionsMethod,
        parallel: bool,
    ) -> Vec<f64> {
        self.boosters
            .iter()
            .flat_map(|b| b.predict_contributions_columnar(data, method, parallel))
            .collect::<Vec<f64>>()
    }

    pub fn predict_intervals(&self, data: &Matrix<f64>, parallel: bool) -> HashMap<String, Vec<Vec<f64>>> {
        let mut results: HashMap<String, Vec<Vec<f64>>> = HashMap::new();
        for booster in &self.boosters {
            let booster_intervals = booster.predict_intervals(data, parallel);
            for (alpha, intervals) in booster_intervals {
                let entry = results.entry(alpha).or_insert_with(|| vec![Vec::new(); data.rows]);
                for (i, sample_interval) in intervals.into_iter().enumerate() {
                    entry[i].extend(sample_interval);
                }
            }
        }
        results
    }

    pub fn predict_intervals_columnar(
        &self,
        data: &ColumnarMatrix<f64>,
        parallel: bool,
    ) -> HashMap<String, Vec<Vec<f64>>> {
        let mut results: HashMap<String, Vec<Vec<f64>>> = HashMap::new();
        let n_samples = data.index.len();
        for booster in &self.boosters {
            let booster_intervals = booster.predict_intervals_columnar(data, parallel);
            for (alpha, intervals) in booster_intervals {
                let entry = results.entry(alpha).or_insert_with(|| vec![Vec::new(); n_samples]);
                for (i, sample_interval) in intervals.into_iter().enumerate() {
                    entry[i].extend(sample_interval);
                }
            }
        }
        results
    }

    /// Predict a distribution using uncalibrated leaf weights from internal nodes.
    /// Returns `n` predictions for each sample, resulting in a matrix of size `(n_samples, n_boosters * n)`.
    pub fn predict_distribution(&self, data: &Matrix<f64>, n: usize, parallel: bool) -> Vec<Vec<f64>> {
        let n_samples = data.rows;
        let mut results = vec![Vec::with_capacity(self.boosters.len() * n); n_samples];
        for booster in self.boosters.iter() {
            let booster_dist = booster.predict_distribution(data, n, parallel);
            for (i, sample_dist) in booster_dist.into_iter().enumerate() {
                results[i].extend(sample_dist);
            }
        }
        results
    }

    pub fn predict_distribution_columnar(&self, data: &ColumnarMatrix<f64>, n: usize, parallel: bool) -> Vec<Vec<f64>> {
        let n_samples = data.index.len();
        let mut results = vec![Vec::with_capacity(self.boosters.len() * n); n_samples];
        for booster in self.boosters.iter() {
            let booster_dist = booster.predict_distribution_columnar(data, n, parallel);
            for (i, sample_dist) in booster_dist.into_iter().enumerate() {
                results[i].extend(sample_dist);
            }
        }
        results
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::booster::core::RegressionLinearHead;
    use crate::node::Node;
    use crate::objective::Objective;
    use crate::tree::core::Tree;
    use approx::assert_relative_eq;

    fn create_mock_booster() -> PerpetualBooster {
        let mut booster = PerpetualBooster::default();
        booster.cfg.objective = Objective::SquaredLoss;
        booster.base_score = 0.5;
        let mut tree = Tree::new();
        let root = Node {
            num: 0,
            weight_value: 0.1,
            leaf_weights: None,
            hessian_sum: 10.0,
            split_value: 0.0,
            split_feature: 0,
            split_gain: 0.0,
            missing_node: 0,
            left_child: 0,
            right_child: 0,
            is_leaf: true,
            parent_node: 0,
            left_cats: None,
            stats: None,
        };
        tree.nodes.insert(0, root);
        booster.trees = vec![tree];
        booster
    }

    fn create_mock_multi_booster() -> MultiOutputBooster {
        let mut first = create_mock_booster();
        first.base_score = 0.5;
        first.trees[0].nodes.get_mut(&0).unwrap().weight_value = 0.1;

        let mut second = create_mock_booster();
        second.base_score = 0.5;
        second.trees[0].nodes.get_mut(&0).unwrap().weight_value = 0.1;

        let cfg = first.cfg.clone();

        MultiOutputBooster {
            n_boosters: 2,
            cfg,
            boosters: vec![first, second],
            class_priors: vec![0.5, 0.5],
            native_multiclass: false,
            metadata: HashMap::new(),
        }
    }

    #[test]
    fn test_predict_basic() {
        let booster = create_mock_booster();
        let data = Matrix::new(&[1.0, 2.0], 1, 2);
        let preds = booster.predict(&data, false);
        assert_eq!(preds.len(), 1);
        assert_relative_eq!(preds[0], 0.6, epsilon = 1e-7);
    }

    #[test]
    fn test_predict_applies_regression_head_for_dense_and_columnar_data() {
        let mut booster = create_mock_booster();
        booster.regression_head = Some(RegressionLinearHead {
            feature_means: vec![0.0, 0.0],
            feature_scales: vec![1.0, 1.0],
            coefficients: vec![0.1, 0.2],
            intercept: 0.5,
            blend_weight: 0.4,
        });

        let data = Matrix::new(&[1.0, 2.0], 1, 2);
        let preds = booster.predict(&data, false);
        assert_relative_eq!(preds[0], 1.0, epsilon = 1e-7);

        let data_vec = [1.0, 2.0];
        let columnar = ColumnarMatrix::new(vec![&data_vec[0..1], &data_vec[1..2]], None, 1);
        let columnar_preds = booster.predict_columnar(&columnar, false);
        assert_relative_eq!(columnar_preds[0], 1.0, epsilon = 1e-7);
    }

    #[test]
    fn test_predict_parallel() {
        let booster = create_mock_booster();
        let data = Matrix::new(&[1.0, 2.0, 3.0, 4.0], 2, 2);
        let preds = booster.predict(&data, true);
        assert_eq!(preds.len(), 2);
        assert_relative_eq!(preds[0], 0.6, epsilon = 1e-7);
    }

    #[test]
    fn test_predict_columnar() {
        let booster = create_mock_booster();
        let data_vec = [1.0, 2.0, 3.0, 4.0];
        let col0 = &data_vec[0..2];
        let col1 = &data_vec[2..4];
        let data = ColumnarMatrix::new(vec![col0, col1], None, 2);
        let preds = booster.predict_columnar(&data, false);
        assert_eq!(preds.len(), 2);
        assert_relative_eq!(preds[0], 0.6, epsilon = 1e-7);
    }

    #[test]
    fn test_predict_proba_basic() {
        let mut booster = create_mock_booster();
        booster.cfg.objective = Objective::LogLoss;
        booster.base_score = 0.0;
        let data = Matrix::new(&[1.0, 2.0], 1, 2);
        let proba = booster.predict_proba(&data, false, false);
        assert_relative_eq!(proba[0], 0.52497918747894, epsilon = 1e-7);
    }

    #[test]
    fn test_predict_distribution() {
        let booster = create_mock_booster();
        let data = Matrix::new(&[1.0, 2.0, 3.0, 4.0], 2, 2);
        let n = 100;
        let dist = booster.predict_distribution(&data, n, false);
        assert_eq!(dist.len(), 2);
        assert_eq!(dist[0].len(), n);

        // base_score (0.5) + leaf weight_value (0.1) = 0.6
        // Since stats is None, it defaults to weight_value repeated 5 times
        for d in dist[0].iter().take(n) {
            assert_relative_eq!(*d, 0.6, epsilon = 1e-7);
        }
        for d in dist[1].iter().take(n) {
            assert_relative_eq!(*d, 0.6, epsilon = 1e-7);
        }
    }

    #[test]
    fn test_predict_proba_parallel() {
        let mut booster = create_mock_booster();
        booster.cfg.objective = Objective::LogLoss;
        booster.base_score = 0.0;
        let data = Matrix::new(&[1.0, 2.0], 1, 2);
        let proba = booster.predict_proba(&data, true, false);
        assert_relative_eq!(proba[0], 0.52497918747894, epsilon = 1e-7);
    }

    #[test]
    fn test_predict_proba_uses_fold_ensemble_when_available() {
        let mut booster = create_mock_booster();
        booster.cfg.objective = Objective::LogLoss;
        booster.base_score = 0.0;
        booster.trees[0].nodes.get_mut(&0).unwrap().leaf_weights = Some([-0.4, -0.1, 0.1, 0.4, 0.8]);

        let data = Matrix::new(&[1.0, 2.0], 1, 2);
        let proba = booster.predict_proba(&data, false, false);
        let expected = [-0.4_f64, -0.1, 0.1, 0.4, 0.8].iter().map(|p| odds(*p)).sum::<f64>() / 5.0;
        assert_relative_eq!(proba[0], expected, epsilon = 1e-7);
    }

    #[test]
    fn test_predict_proba_calibrated_uses_same_fold_ensemble_before_calibration() {
        let mut booster = create_mock_booster();
        booster.cfg.objective = Objective::LogLoss;
        booster.base_score = 0.0;
        booster.trees[0].nodes.get_mut(&0).unwrap().leaf_weights = Some([-0.4, -0.1, 0.1, 0.4, 0.8]);

        let data = Matrix::new(&[1.0, 2.0], 1, 2);
        let proba = booster.predict_proba(&data, false, true);
        let expected = [-0.4_f64, -0.1, 0.1, 0.4, 0.8].iter().map(|p| odds(*p)).sum::<f64>() / 5.0;
        assert_relative_eq!(proba[0], expected, epsilon = 1e-7);
    }

    #[test]
    fn test_predict_contributions_methods() {
        let booster = create_mock_booster();
        let data = Matrix::new(&[1.0, 2.0], 1, 2);
        let methods = vec![
            ContributionsMethod::Weight,
            ContributionsMethod::Average,
            ContributionsMethod::MidpointDifference,
            ContributionsMethod::BranchDifference,
            ContributionsMethod::ModeDifference,
            ContributionsMethod::Shapley,
        ];
        for method in methods {
            let contribs = booster.predict_contributions(&data, method, false);
            assert_eq!(contribs.len(), 3);
        }
    }

    #[test]
    fn test_predict_contributions_probability_change() {
        let mut booster = create_mock_booster();
        booster.cfg.objective = Objective::LogLoss;
        let data = Matrix::new(&[1.0, 2.0], 1, 2);
        let contribs = booster.predict_contributions(&data, ContributionsMethod::ProbabilityChange, false);
        assert_eq!(contribs.len(), 3);
        // Sum of contributions + bias should ≈ probability?
        // Actually for ProbabilityChange it depends on the implementation.
    }

    #[test]
    fn test_predict_proba_calibrated() {
        let mut booster = create_mock_booster();
        booster.cfg.objective = Objective::LogLoss;
        // Adding a dummy calibrator would be complex, but let's see if we can trigger the code path.
        let data = Matrix::new(&[1.0, 2.0], 1, 2);
        let proba = booster.predict_proba(&data, false, true);
        assert_eq!(proba.len(), 1);
    }

    #[test]
    fn test_predict_nodes() {
        let booster = create_mock_booster();
        let data = Matrix::new(&[1.0, 2.0], 1, 2);
        let nodes = booster.predict_nodes(&data, false);
        assert_eq!(nodes.len(), 1);
        assert_eq!(nodes[0].len(), 1);
        assert!(nodes[0][0].contains(&0));
    }

    #[test]
    fn test_multi_output_predict_basic() {
        let booster = create_mock_multi_booster();
        let data = Matrix::new(&[1.0, 2.0], 1, 2);
        let preds = booster.predict(&data, false);
        assert_eq!(preds.len(), 2);
        assert_relative_eq!(preds[0], 0.6, epsilon = 1e-7);
    }

    #[test]
    fn test_multi_output_predict_proba_basic() {
        let booster = create_mock_multi_booster();
        let data = Matrix::new(&[1.0, 2.0], 1, 2);
        let proba = booster.predict_proba(&data, false);
        assert_eq!(proba.len(), 2);
        assert_relative_eq!(proba[0], 0.5, epsilon = 1e-7);
    }

    #[test]
    fn test_multi_output_predict_proba_normalizes_binary_probabilities() {
        let mut first = create_mock_booster();
        first.cfg.objective = Objective::LogLoss;
        first.base_score = 0.0;
        first.trees.clear();

        let mut second = create_mock_booster();
        second.cfg.objective = Objective::LogLoss;
        second.base_score = 2.0;
        second.trees.clear();

        let booster = MultiOutputBooster {
            n_boosters: 2,
            boosters: vec![first, second],
            class_priors: vec![0.5, 0.5],
            ..Default::default()
        };
        let data = Matrix::new(&[1.0, 2.0], 1, 2);
        let proba = booster.predict_proba(&data, false);

        assert_eq!(proba.len(), 2);
        assert_relative_eq!(proba[0], 0.362_109_7, epsilon = 1e-6);
        assert_relative_eq!(proba[1], 0.637_890_3, epsilon = 1e-6);
    }

    #[test]
    fn test_multi_output_predict_intervals_columnar() {
        let booster = create_mock_multi_booster();
        let data_vec = [1.0, 2.0, 3.0, 4.0];
        let col0 = &data_vec[0..2];
        let col1 = &data_vec[2..4];
        let data = ColumnarMatrix::new(vec![col0, col1], None, 2);
        let intervals = booster.predict_intervals_columnar(&data, false);
        assert!(intervals.is_empty());
    }

    #[test]
    fn test_predict_contributions_columnar() {
        let booster = create_mock_booster();
        let data_vec = [1.0, 2.0, 3.0, 4.0];
        let col0 = &data_vec[0..2];
        let col1 = &data_vec[2..4];
        let data = ColumnarMatrix::new(vec![col0, col1], None, 2);
        let contribs = booster.predict_contributions_columnar(&data, ContributionsMethod::Weight, false);
        assert_eq!(contribs.len(), 6); // 2 rows * (2 features + 1 bias)
    }

    #[test]
    fn test_multi_output_predict_contributions_columnar() {
        let booster = create_mock_multi_booster();
        let data_vec = [1.0, 2.0, 3.0, 4.0];
        let col0 = &data_vec[0..2];
        let col1 = &data_vec[2..4];
        let data = ColumnarMatrix::new(vec![col0, col1], None, 2);
        let contribs = booster.predict_contributions_columnar(&data, ContributionsMethod::Weight, false);
        assert_eq!(contribs.len(), 12); // 2 rows * 3 items * 2 boosters
    }

    // ---- Additional tests for uncovered paths ----

    #[test]
    fn test_predict_proba_columnar() {
        let mut booster = create_mock_booster();
        booster.cfg.objective = Objective::LogLoss;
        booster.base_score = 0.0;
        let data_vec = [1.0, 2.0, 3.0, 4.0];
        let col0 = &data_vec[0..2];
        let col1 = &data_vec[2..4];
        let data = ColumnarMatrix::new(vec![col0, col1], None, 2);
        let proba = booster.predict_proba_columnar(&data, false, false);
        assert_eq!(proba.len(), 2);
        for p in &proba {
            assert!(*p > 0.0 && *p < 1.0);
        }
    }

    #[test]
    fn test_predict_proba_columnar_parallel() {
        let mut booster = create_mock_booster();
        booster.cfg.objective = Objective::LogLoss;
        booster.base_score = 0.0;
        let data_vec = [1.0, 2.0, 3.0, 4.0];
        let col0 = &data_vec[0..2];
        let col1 = &data_vec[2..4];
        let data = ColumnarMatrix::new(vec![col0, col1], None, 2);
        let proba = booster.predict_proba_columnar(&data, true, false);
        assert_eq!(proba.len(), 2);
    }

    #[test]
    fn test_predict_nodes_columnar() {
        let booster = create_mock_booster();
        let data_vec = [1.0, 2.0, 3.0, 4.0];
        let col0 = &data_vec[0..2];
        let col1 = &data_vec[2..4];
        let data = ColumnarMatrix::new(vec![col0, col1], None, 2);
        let nodes = booster.predict_nodes_columnar(&data, false);
        assert_eq!(nodes.len(), 1); // 1 tree
        assert_eq!(nodes[0].len(), 2); // 2 rows
    }

    #[test]
    fn test_predict_contributions_average_parallel() {
        let booster = create_mock_booster();
        let data = Matrix::new(&[1.0, 2.0, 3.0, 4.0], 2, 2);
        let contribs = booster.predict_contributions(&data, ContributionsMethod::Average, true);
        assert_eq!(contribs.len(), 6); // 2 rows * (2 features + 1 bias)
    }

    #[test]
    fn test_predict_contributions_weight_parallel() {
        let booster = create_mock_booster();
        let data = Matrix::new(&[1.0, 2.0, 3.0, 4.0], 2, 2);
        let contribs = booster.predict_contributions(&data, ContributionsMethod::Weight, true);
        assert_eq!(contribs.len(), 6);
    }

    #[test]
    fn test_predict_contributions_columnar_average() {
        let booster = create_mock_booster();
        let data_vec = [1.0, 2.0, 3.0, 4.0];
        let col0 = &data_vec[0..2];
        let col1 = &data_vec[2..4];
        let data = ColumnarMatrix::new(vec![col0, col1], None, 2);
        let contribs = booster.predict_contributions_columnar(&data, ContributionsMethod::Average, false);
        assert_eq!(contribs.len(), 6);
    }

    #[test]
    fn test_predict_contributions_columnar_average_parallel() {
        let booster = create_mock_booster();
        let data_vec = [1.0, 2.0, 3.0, 4.0];
        let col0 = &data_vec[0..2];
        let col1 = &data_vec[2..4];
        let data = ColumnarMatrix::new(vec![col0, col1], None, 2);
        let contribs = booster.predict_contributions_columnar(&data, ContributionsMethod::Average, true);
        assert_eq!(contribs.len(), 6);
    }

    #[test]
    fn test_predict_contributions_columnar_all_methods() {
        let booster = create_mock_booster();
        let data_vec = [1.0, 2.0, 3.0, 4.0];
        let col0 = &data_vec[0..2];
        let col1 = &data_vec[2..4];
        let data = ColumnarMatrix::new(vec![col0, col1], None, 2);
        let methods = vec![
            ContributionsMethod::Weight,
            ContributionsMethod::BranchDifference,
            ContributionsMethod::MidpointDifference,
            ContributionsMethod::ModeDifference,
            ContributionsMethod::Shapley,
        ];
        for method in methods {
            let contribs = booster.predict_contributions_columnar(&data, method, false);
            assert_eq!(contribs.len(), 6);
        }
    }

    #[test]
    fn test_predict_contributions_columnar_probability_change() {
        let mut booster = create_mock_booster();
        booster.cfg.objective = Objective::LogLoss;
        let data_vec = [1.0, 2.0, 3.0, 4.0];
        let col0 = &data_vec[0..2];
        let col1 = &data_vec[2..4];
        let data = ColumnarMatrix::new(vec![col0, col1], None, 2);
        let contribs = booster.predict_contributions_columnar(&data, ContributionsMethod::ProbabilityChange, false);
        assert_eq!(contribs.len(), 6);
    }

    #[test]
    fn test_predict_contributions_columnar_probability_change_parallel() {
        let mut booster = create_mock_booster();
        booster.cfg.objective = Objective::LogLoss;
        let data_vec = [1.0, 2.0, 3.0, 4.0];
        let col0 = &data_vec[0..2];
        let col1 = &data_vec[2..4];
        let data = ColumnarMatrix::new(vec![col0, col1], None, 2);
        let contribs = booster.predict_contributions_columnar(&data, ContributionsMethod::ProbabilityChange, true);
        assert_eq!(contribs.len(), 6);
    }

    #[test]
    fn test_predict_contributions_tree_alone_columnar_parallel() {
        let booster = create_mock_booster();
        let data_vec = [1.0, 2.0, 3.0, 4.0];
        let col0 = &data_vec[0..2];
        let col1 = &data_vec[2..4];
        let data = ColumnarMatrix::new(vec![col0, col1], None, 2);
        let contribs = booster.predict_contributions_columnar(&data, ContributionsMethod::Weight, true);
        assert_eq!(contribs.len(), 6);
    }

    #[test]
    fn test_multi_output_predict_columnar() {
        let booster = create_mock_multi_booster();
        let data_vec = [1.0, 2.0, 3.0, 4.0];
        let col0 = &data_vec[0..2];
        let col1 = &data_vec[2..4];
        let data = ColumnarMatrix::new(vec![col0, col1], None, 2);
        let preds = booster.predict_columnar(&data, false);
        assert_eq!(preds.len(), 4); // 2 rows * 2 boosters
    }

    #[test]
    fn test_multi_output_predict_proba_columnar() {
        let booster = create_mock_multi_booster();
        let data_vec = [1.0, 2.0, 3.0, 4.0];
        let col0 = &data_vec[0..2];
        let col1 = &data_vec[2..4];
        let data = ColumnarMatrix::new(vec![col0, col1], None, 2);
        let proba = booster.predict_proba_columnar(&data, false);
        assert_eq!(proba.len(), 4);
        // Probabilities should sum to 1 per row
    }

    #[test]
    fn test_multi_output_predict_proba_columnar_normalizes_binary_probabilities() {
        let mut first = create_mock_booster();
        first.cfg.objective = Objective::LogLoss;
        first.base_score = 0.0;
        first.trees.clear();

        let mut second = create_mock_booster();
        second.cfg.objective = Objective::LogLoss;
        second.base_score = 2.0;
        second.trees.clear();

        let booster = MultiOutputBooster {
            n_boosters: 2,
            boosters: vec![first, second],
            class_priors: vec![0.5, 0.5],
            ..Default::default()
        };
        let data_vec = [1.0, 2.0, 3.0, 4.0];
        let col0 = &data_vec[0..2];
        let col1 = &data_vec[2..4];
        let data = ColumnarMatrix::new(vec![col0, col1], None, 2);
        let proba = booster.predict_proba_columnar(&data, false);

        assert_eq!(proba.len(), 4);
        assert_relative_eq!(proba[0], 0.362_109_7, epsilon = 1e-6);
        assert_relative_eq!(proba[1], 0.637_890_3, epsilon = 1e-6);
    }

    #[test]
    fn test_multiclass_probability_alpha_prefers_small_class_odds_coupling() {
        let booster = MultiOutputBooster {
            n_boosters: 3,
            class_priors: vec![0.5, 0.3, 0.2],
            ..Default::default()
        };
        let mut probabilities = vec![0.5, 0.880_797_077_977_882_3, 0.2];
        MultiOutputBooster::normalize_multiclass_probabilities(
            &mut probabilities,
            booster.multiclass_probability_alpha(),
            Some(&booster.class_priors),
            booster.multiclass_probability_beta(),
        );

        assert!(probabilities[1] > 0.55);
        assert!(probabilities[2] < 0.12);
    }

    #[test]
    fn test_multiclass_probability_beta_boosts_rare_three_class_prior() {
        let booster = MultiOutputBooster {
            n_boosters: 3,
            class_priors: vec![0.7, 0.2, 0.1],
            ..Default::default()
        };
        let mut with_prior_coupling = vec![0.35, 0.35, 0.35];
        let mut without_prior_coupling = with_prior_coupling.clone();

        MultiOutputBooster::normalize_multiclass_probabilities(
            &mut with_prior_coupling,
            booster.multiclass_probability_alpha(),
            Some(&booster.class_priors),
            booster.multiclass_probability_beta(),
        );
        MultiOutputBooster::normalize_multiclass_probabilities(
            &mut without_prior_coupling,
            booster.multiclass_probability_alpha(),
            Some(&booster.class_priors),
            0.0,
        );

        assert!(with_prior_coupling[2] > without_prior_coupling[2]);
        assert!(with_prior_coupling[0] < without_prior_coupling[0]);
    }

    #[test]
    fn test_native_multiclass_temperature_scales_with_class_count() {
        let native_five = MultiOutputBooster {
            n_boosters: 5,
            native_multiclass: true,
            ..Default::default()
        };
        let native_eight = MultiOutputBooster {
            n_boosters: 8,
            native_multiclass: true,
            ..Default::default()
        };
        let ovr_three = MultiOutputBooster {
            n_boosters: 3,
            native_multiclass: false,
            ..Default::default()
        };

        assert_relative_eq!(native_five.native_multiclass_temperature(), 1.5, epsilon = 1e-12);
        assert_relative_eq!(native_eight.native_multiclass_temperature(), 2.0, epsilon = 1e-12);
        assert_relative_eq!(ovr_three.native_multiclass_temperature(), 1.0, epsilon = 1e-12);
    }

    #[test]
    fn test_native_multiclass_predict_proba_applies_temperature_scaling() {
        let logits = [3.0_f64, 1.0, 0.0, -1.0, -2.0];
        let boosters = logits
            .iter()
            .map(|&base_score| {
                let mut booster = create_mock_booster();
                booster.base_score = base_score;
                booster.trees.clear();
                booster
            })
            .collect::<Vec<_>>();
        let booster = MultiOutputBooster {
            n_boosters: logits.len(),
            boosters,
            native_multiclass: true,
            ..Default::default()
        };
        let data = Matrix::new(&[1.0, 2.0], 1, 2);

        let proba = booster.predict_proba(&data, false);
        let expected = MultiOutputBooster::softmax_probabilities_with_temperature(&logits, 1.5);

        assert_eq!(proba.len(), logits.len());
        proba
            .iter()
            .zip(expected.iter())
            .for_each(|(actual, expected_value)| assert_relative_eq!(*actual, *expected_value, epsilon = 1e-7));
    }

    #[test]
    fn test_multi_output_predict_nodes() {
        let booster = create_mock_multi_booster();
        let data = Matrix::new(&[1.0, 2.0], 1, 2);
        let nodes = booster.predict_nodes(&data, false);
        assert_eq!(nodes.len(), 2); // 2 boosters
    }

    #[test]
    fn test_multi_output_predict_nodes_columnar() {
        let booster = create_mock_multi_booster();
        let data_vec = [1.0, 2.0, 3.0, 4.0];
        let col0 = &data_vec[0..2];
        let col1 = &data_vec[2..4];
        let data = ColumnarMatrix::new(vec![col0, col1], None, 2);
        let nodes = booster.predict_nodes_columnar(&data, false);
        assert_eq!(nodes.len(), 2); // 2 boosters
    }

    #[test]
    fn test_multi_output_predict_contributions() {
        let booster = create_mock_multi_booster();
        let data = Matrix::new(&[1.0, 2.0], 1, 2);
        let contribs = booster.predict_contributions(&data, ContributionsMethod::Weight, false);
        assert_eq!(contribs.len(), 6); // 1 row * 3 items * 2 boosters
    }

    #[test]
    fn test_multi_output_predict_intervals() {
        let booster = create_mock_multi_booster();
        let data = Matrix::new(&[1.0, 2.0], 1, 2);
        let intervals = booster.predict_intervals(&data, false);
        // No calibration params set, so intervals should be empty
        assert!(intervals.is_empty());
    }

    #[test]
    fn test_predict_contributions_probability_change_parallel() {
        let mut booster = create_mock_booster();
        booster.cfg.objective = Objective::LogLoss;
        let data = Matrix::new(&[1.0, 2.0, 3.0, 4.0], 2, 2);
        let contribs = booster.predict_contributions(&data, ContributionsMethod::ProbabilityChange, true);
        assert_eq!(contribs.len(), 6);
    }
}