irithyll 10.0.1

//! Unified streaming learner trait for polymorphic model composition.
//!
//! [`StreamingLearner`] is an **object-safe** trait that abstracts over any
//! online/streaming machine learning model -- gradient boosted trees, linear
//! models, Naive Bayes, Mondrian forests, or anything else that can ingest
//! samples one at a time and produce predictions.
//!
//! # Motivation
//!
//! Stacking ensembles and meta-learners need to treat heterogeneous base
//! models uniformly: train them on the same stream, collect their predictions
//! as features for a combiner, and manage their lifecycle (reset, clone,
//! serialization). `StreamingLearner` provides exactly this interface.
//!
//! # Object Safety
//!
//! The trait is deliberately object-safe: every method uses `&self` /
//! `&mut self` with concrete return types (no generics on methods, no
//! `Self`-by-value in non-`Sized` positions). This means you can store
//! `Box<dyn StreamingLearner>` in a `Vec`, enabling runtime-polymorphic
//! stacking without monomorphization.
//!
//! # Usage
//!
//! ```
//! use irithyll::learner::{StreamingLearner, SGBTLearner};
//! use irithyll::SGBTConfig;
//!
//! // Create a base learner from config
//! let config = SGBTConfig::builder()
//!     .n_steps(10)
//!     .learning_rate(0.1)
//!     .build()
//!     .unwrap();
//! let mut learner = SGBTLearner::from_config(config);
//!
//! // Train incrementally
//! learner.train(&[1.0, 2.0], 3.0);
//! learner.train(&[4.0, 5.0], 6.0);
//!
//! // Predict
//! let pred = learner.predict(&[1.0, 2.0]);
//! assert!(pred.is_finite());
//!
//! // Use as trait object for stacking
//! let boxed: Box<dyn StreamingLearner> = Box::new(learner);
//! assert_eq!(boxed.n_samples_seen(), 2);
//! ```

use std::fmt;

use crate::ensemble::config::SGBTConfig;
use crate::ensemble::SGBT;
use crate::loss::squared::SquaredLoss;
use crate::loss::Loss;
use crate::sample::SampleRef;

// Re-export the trait and capability traits from irithyll-core so downstream
// code can use `irithyll::learner::StreamingLearner` unchanged and access the
// new capability traits from the same module path.
pub use irithyll_core::learner::{HasReadout, StreamingLearner, Structural, Tunable};

// ---------------------------------------------------------------------------
// SGBTLearner -- adapter wrapping SGBT<L> into StreamingLearner
// ---------------------------------------------------------------------------

/// Adapter that wraps an [`SGBT`] ensemble into the [`StreamingLearner`] trait.
///
/// This is the primary way to use SGBT models in polymorphic stacking
/// ensembles. The loss function `L` is monomorphized at compile time for
/// zero-cost gradient dispatch, while the `StreamingLearner` trait provides
/// the uniform interface needed by meta-learners.
///
/// # Type Parameter
///
/// * `L` -- loss function type, defaulting to [`SquaredLoss`] for regression.
///   Any `L: Loss + Clone` is supported.
///
/// # Examples
///
/// ```
/// use irithyll::learner::SGBTLearner;
/// use irithyll::SGBTConfig;
///
/// // Default squared loss:
/// let config = SGBTConfig::builder().n_steps(5).build().unwrap();
/// let learner = SGBTLearner::from_config(config);
/// ```
///
/// ```
/// use irithyll::learner::SGBTLearner;
/// use irithyll::{SGBTConfig, SGBT};
/// use irithyll::loss::logistic::LogisticLoss;
///
/// // Custom loss via wrapping an existing SGBT:
/// let config = SGBTConfig::builder().n_steps(5).build().unwrap();
/// let model = SGBT::with_loss(config, LogisticLoss);
/// let learner = SGBTLearner::new(model);
/// ```
pub struct SGBTLearner<L: Loss = SquaredLoss> {
    inner: SGBT<L>,
}

impl<L: Loss> SGBTLearner<L> {
    /// Wrap an existing [`SGBT`] model into a `SGBTLearner`.
    ///
    /// The model retains all of its current state (trained trees, samples
    /// seen, etc.). This enables wrapping a partially-trained model.
    #[inline]
    pub fn new(model: SGBT<L>) -> Self {
        Self { inner: model }
    }

    /// Immutable access to the underlying [`SGBT`] model.
    ///
    /// Useful for inspecting model internals (config, base prediction,
    /// number of steps) without consuming the adapter.
    #[inline]
    pub fn inner(&self) -> &SGBT<L> {
        &self.inner
    }

    /// Mutable access to the underlying [`SGBT`] model.
    ///
    /// Enables calling SGBT-specific methods (e.g., `predict_transformed`,
    /// serialization) that are not part of the `StreamingLearner` interface.
    #[inline]
    pub fn inner_mut(&mut self) -> &mut SGBT<L> {
        &mut self.inner
    }

    /// Consume the adapter and return the underlying [`SGBT`] model.
    ///
    /// This is useful when you need to serialize the model or switch from
    /// polymorphic back to monomorphic usage.
    #[inline]
    pub fn into_inner(self) -> SGBT<L> {
        self.inner
    }
}

impl SGBTLearner<SquaredLoss> {
    /// Create a new `SGBTLearner` with squared loss from a configuration.
    ///
    /// This is the most common constructor for regression tasks. For custom
    /// losses, construct the [`SGBT`] first with [`SGBT::with_loss`] and
    /// wrap it via [`SGBTLearner::new`].
    ///
    /// # Examples
    ///
    /// ```
    /// use irithyll::learner::{SGBTLearner, StreamingLearner};
    /// use irithyll::SGBTConfig;
    ///
    /// let config = SGBTConfig::builder()
    ///     .n_steps(10)
    ///     .learning_rate(0.05)
    ///     .build()
    ///     .unwrap();
    /// let learner = SGBTLearner::from_config(config);
    /// assert_eq!(learner.n_samples_seen(), 0);
    /// ```
    #[inline]
    pub fn from_config(config: SGBTConfig) -> Self {
        Self {
            inner: SGBT::new(config),
        }
    }
}

// ---------------------------------------------------------------------------
// StreamingLearner impl for SGBTLearner
// ---------------------------------------------------------------------------

impl<L: Loss> StreamingLearner for SGBTLearner<L> {
    #[inline]
    fn train_one(&mut self, features: &[f64], target: f64, weight: f64) {
        let sample = SampleRef::weighted(features, target, weight);
        self.inner.train_one(&sample);
    }

    #[inline]
    fn predict(&self, features: &[f64]) -> f64 {
        self.inner.predict(features)
    }

    #[inline]
    fn n_samples_seen(&self) -> u64 {
        self.inner.n_samples_seen()
    }

    #[inline]
    fn reset(&mut self) {
        self.inner.reset();
    }

    #[allow(deprecated)]
    fn diagnostics_array(&self) -> [f64; 5] {
        <Self as Tunable>::diagnostics_array(self)
    }

    #[allow(deprecated)]
    fn adjust_config(&mut self, lr_multiplier: f64, lambda_delta: f64) {
        <Self as Tunable>::adjust_config(self, lr_multiplier, lambda_delta);
    }

    #[allow(deprecated)]
    fn apply_structural_change(&mut self, depth_delta: i32, steps_delta: i32) {
        <Self as Structural>::apply_structural_change(self, depth_delta, steps_delta);
    }

    #[allow(deprecated)]
    fn replacement_count(&self) -> u64 {
        <Self as Structural>::replacement_count(self)
    }

    #[allow(deprecated)]
    fn check_proactive_prune(&mut self) -> bool {
        <Self as Structural>::check_proactive_prune(self)
    }

    #[allow(deprecated)]
    fn set_prune_half_life(&mut self, hl: usize) {
        <Self as Structural>::set_prune_half_life(self, hl);
    }

    #[allow(deprecated)]
    fn tree_structure(&self) -> Vec<(usize, usize, f64, f64, u64)> {
        <Self as Structural>::tree_structure(self)
    }
}

// ---------------------------------------------------------------------------
// Tunable impl -- exposes diagnostics and config-adjust to AutoML components
// ---------------------------------------------------------------------------

impl<L: Loss> Tunable for SGBTLearner<L> {
    fn diagnostics_array(&self) -> [f64; 5] {
        use crate::automl::DiagnosticSource;
        match self.inner.config_diagnostics() {
            Some(d) => [
                d.residual_alignment,
                d.regularization_sensitivity,
                d.depth_sufficiency,
                d.effective_dof,
                d.uncertainty,
            ],
            None => [0.0; 5],
        }
    }

    fn adjust_config(&mut self, lr_multiplier: f64, lambda_delta: f64) {
        let current_lr = self.inner.config().learning_rate;
        self.inner.set_learning_rate(current_lr * lr_multiplier);
        let current_lambda = self.inner.config().lambda;
        self.inner.set_lambda(current_lambda + lambda_delta);
    }
}

// ---------------------------------------------------------------------------
// Structural impl -- exposes tree capacity management to AutoML components
// ---------------------------------------------------------------------------

impl<L: Loss> Structural for SGBTLearner<L> {
    fn apply_structural_change(&mut self, depth_delta: i32, steps_delta: i32) {
        if depth_delta != 0 {
            let current = self.inner.config().max_depth as i32;
            self.inner
                .set_max_depth((current + depth_delta).max(1) as usize);
        }
        if steps_delta != 0 {
            let current = self.inner.config().n_steps as i32;
            self.inner
                .set_n_steps((current + steps_delta).max(3) as usize);
        }
    }

    fn replacement_count(&self) -> u64 {
        self.inner.total_replacements()
    }

    fn check_proactive_prune(&mut self) -> bool {
        self.inner.check_proactive_prune()
    }

    fn set_prune_half_life(&mut self, hl: usize) {
        self.inner.set_prune_half_life(hl);
    }

    fn tree_structure(&self) -> Vec<(usize, usize, f64, f64, u64)> {
        let diag = self.inner.diagnostics_overview();
        diag.trees
            .iter()
            .map(|t| (t.max_depth, t.n_leaves, 0.0, 0.0, t.n_samples))
            .collect()
    }
}

// ---------------------------------------------------------------------------
// Clone impl -- manual to match irithyll patterns (no derive)
// ---------------------------------------------------------------------------

impl<L: Loss + Clone> Clone for SGBTLearner<L> {
    fn clone(&self) -> Self {
        Self {
            inner: self.inner.clone(),
        }
    }
}

// ---------------------------------------------------------------------------
// Debug impl -- manual, avoids requiring Debug on L
// ---------------------------------------------------------------------------

impl<L: Loss> fmt::Debug for SGBTLearner<L> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("SGBTLearner")
            .field("inner", &self.inner)
            .finish()
    }
}

// ---------------------------------------------------------------------------
// DiagnosticSource impl
// ---------------------------------------------------------------------------

impl<L: Loss> crate::automl::DiagnosticSource for SGBTLearner<L> {
    fn config_diagnostics(&self) -> Option<crate::automl::ConfigDiagnostics> {
        // Forward to the inner SGBT's diagnostics (which has the real signals).
        self.inner.config_diagnostics()
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::SGBTConfig;

    /// Shared minimal config for tests.
    fn default_config() -> SGBTConfig {
        SGBTConfig::builder()
            .n_steps(10)
            .learning_rate(0.1)
            .grace_period(20)
            .max_depth(4)
            .n_bins(16)
            .build()
            .unwrap()
    }

    #[test]
    fn test_sgbt_learner_creation() {
        let learner = SGBTLearner::from_config(default_config());
        assert_eq!(learner.n_samples_seen(), 0);
        // Untrained model predicts zero (no base prediction initialized).
        let pred = learner.predict(&[1.0, 2.0, 3.0]);
        assert!(pred.abs() < 1e-12);
    }

    #[test]
    fn test_trait_object_safety() {
        // Verify that Box<dyn StreamingLearner> compiles and works.
        let learner = SGBTLearner::from_config(default_config());
        let mut boxed: Box<dyn StreamingLearner> = Box::new(learner);

        // Train through the trait object.
        boxed.train(&[1.0, 2.0, 3.0], 5.0);
        assert_eq!(boxed.n_samples_seen(), 1);

        // Predict through the trait object.
        let pred = boxed.predict(&[1.0, 2.0, 3.0]);
        assert!(pred.is_finite());

        // Reset through the trait object.
        boxed.reset();
        assert_eq!(boxed.n_samples_seen(), 0);
    }

    #[test]
    fn test_train_and_predict() {
        let mut learner = SGBTLearner::from_config(default_config());

        // Record prediction before training.
        let features = [1.0, 2.0, 3.0];
        let pred_before = learner.predict(&features);

        // Train on enough samples to initialize base prediction and grow trees.
        for i in 0..100 {
            learner.train(&features, (i as f64) * 0.1);
        }
        assert_eq!(learner.n_samples_seen(), 100);

        // Prediction should have changed from the untrained state.
        let pred_after = learner.predict(&features);
        assert!(
            (pred_after - pred_before).abs() > 1e-6,
            "prediction should change after training: before={}, after={}",
            pred_before,
            pred_after,
        );
        assert!(pred_after.is_finite());
    }

    #[test]
    fn test_reset() {
        let mut learner = SGBTLearner::from_config(default_config());

        // Train some samples.
        for i in 0..50 {
            learner.train(&[1.0, 2.0], i as f64);
        }
        assert_eq!(learner.n_samples_seen(), 50);

        // Reset and verify state is cleared.
        learner.reset();
        assert_eq!(learner.n_samples_seen(), 0);

        // After reset, predictions should return to zero (uninitialised base).
        let pred = learner.predict(&[1.0, 2.0]);
        assert!(
            pred.abs() < 1e-12,
            "prediction after reset should be zero, got {}",
            pred,
        );
    }

    #[test]
    fn test_predict_batch() {
        let mut learner = SGBTLearner::from_config(default_config());

        // Train a bit so predictions are non-trivial.
        for i in 0..60 {
            learner.train(&[i as f64, (i as f64) * 0.5], i as f64);
        }

        let rows: Vec<&[f64]> = vec![&[1.0, 0.5], &[10.0, 5.0], &[30.0, 15.0]];
        let batch = learner.predict_batch(&rows);

        // Batch results should match individual predictions exactly.
        assert_eq!(batch.len(), rows.len());
        for (i, row) in rows.iter().enumerate() {
            let individual = learner.predict(row);
            assert!(
                (batch[i] - individual).abs() < 1e-12,
                "batch[{}]={} != individual={}",
                i,
                batch[i],
                individual,
            );
        }
    }

    #[test]
    fn test_inner_access() {
        let config = default_config();
        let learner = SGBTLearner::from_config(config.clone());

        // inner() should return the SGBT reference with matching config.
        let sgbt = learner.inner();
        assert_eq!(sgbt.n_samples_seen(), 0);
        assert_eq!(sgbt.config().n_steps, config.n_steps);
    }

    #[test]
    fn test_clone() {
        let mut learner = SGBTLearner::from_config(default_config());

        // Train the original.
        for i in 0..80 {
            learner.train(&[i as f64, (i as f64) * 2.0], i as f64);
        }

        // Clone should be an independent copy.
        let mut cloned = learner.clone();
        assert_eq!(cloned.n_samples_seen(), learner.n_samples_seen());

        // Predictions should match immediately after clone.
        let features = [5.0, 10.0];
        let pred_orig = learner.predict(&features);
        let pred_clone = cloned.predict(&features);
        assert!(
            (pred_orig - pred_clone).abs() < 1e-12,
            "clone prediction should match original: {} vs {}",
            pred_orig,
            pred_clone,
        );

        // Training the clone should not affect the original.
        for i in 0..50 {
            cloned.train(&[i as f64, (i as f64) * 2.0], 999.0);
        }
        assert_eq!(learner.n_samples_seen(), 80);
        assert_eq!(cloned.n_samples_seen(), 130);

        // Predictions should now diverge.
        let pred_orig_after = learner.predict(&features);
        let pred_clone_after = cloned.predict(&features);
        assert!(
            (pred_orig - pred_orig_after).abs() < 1e-12,
            "original should be unchanged after training clone",
        );
        assert!(
            (pred_clone_after - pred_clone).abs() > 1e-6,
            "clone prediction should change after further training",
        );
    }
}