irithyll 10.0.1

//! Model serialization and deserialization support.
//!
//! Provides JSON serialization (default feature) and optional bincode
//! serialization for model persistence. Currently exposes utility functions
//! for serializing/deserializing generic `Serialize`/`Deserialize` types.

use crate::error::{IrithyllError, Result};
use serde::{Deserialize, Serialize};

/// Serialize a value to a JSON string.
///
/// Requires the `serde-json` feature (enabled by default).
///
/// # Errors
///
/// Returns [`IrithyllError::Serialization`] if serialization fails.
#[cfg(feature = "serde-json")]
#[cfg_attr(docsrs, doc(cfg(feature = "serde-json")))]
pub fn to_json<T: serde::Serialize>(value: &T) -> Result<String> {
    serde_json::to_string(value).map_err(|e| IrithyllError::Serialization(e.to_string()))
}

/// Serialize a value to a pretty-printed JSON string.
///
/// Requires the `serde-json` feature (enabled by default).
///
/// # Errors
///
/// Returns [`IrithyllError::Serialization`] if serialization fails.
#[cfg(feature = "serde-json")]
#[cfg_attr(docsrs, doc(cfg(feature = "serde-json")))]
pub fn to_json_pretty<T: serde::Serialize>(value: &T) -> Result<String> {
    serde_json::to_string_pretty(value).map_err(|e| IrithyllError::Serialization(e.to_string()))
}

/// Deserialize a value from a JSON string.
///
/// Requires the `serde-json` feature (enabled by default).
///
/// # Errors
///
/// Returns [`IrithyllError::Serialization`] if deserialization fails.
#[cfg(feature = "serde-json")]
#[cfg_attr(docsrs, doc(cfg(feature = "serde-json")))]
pub fn from_json<T: serde::de::DeserializeOwned>(json: &str) -> Result<T> {
    serde_json::from_str(json).map_err(|e| IrithyllError::Serialization(e.to_string()))
}

/// Serialize a value to JSON bytes.
///
/// Requires the `serde-json` feature (enabled by default).
///
/// # Errors
///
/// Returns [`IrithyllError::Serialization`] if serialization fails.
#[cfg(feature = "serde-json")]
#[cfg_attr(docsrs, doc(cfg(feature = "serde-json")))]
pub fn to_json_bytes<T: serde::Serialize>(value: &T) -> Result<Vec<u8>> {
    serde_json::to_vec(value).map_err(|e| IrithyllError::Serialization(e.to_string()))
}

/// Deserialize a value from JSON bytes.
///
/// Requires the `serde-json` feature (enabled by default).
///
/// # Errors
///
/// Returns [`IrithyllError::Serialization`] if deserialization fails.
#[cfg(feature = "serde-json")]
#[cfg_attr(docsrs, doc(cfg(feature = "serde-json")))]
pub fn from_json_bytes<T: serde::de::DeserializeOwned>(bytes: &[u8]) -> Result<T> {
    serde_json::from_slice(bytes).map_err(|e| IrithyllError::Serialization(e.to_string()))
}

// ---------------------------------------------------------------------------
// Model checkpoint/restore serialization
// ---------------------------------------------------------------------------

#[cfg(any(feature = "serde-json", feature = "serde-bincode"))]
use crate::ensemble::config::SGBTConfig;
#[cfg(any(feature = "serde-json", feature = "serde-bincode"))]
use crate::ensemble::SGBT;
#[cfg(any(feature = "serde-json", feature = "serde-bincode"))]
use crate::loss::Loss;

// Re-export LossType from the loss module for backwards compatibility.
#[cfg(any(feature = "serde-json", feature = "serde-bincode"))]
pub use crate::loss::LossType;

/// Serializable snapshot of the tree arena structure.
///
/// Captures the minimal state needed to reconstruct a tree for prediction:
/// node topology, split decisions, and leaf values. Histogram accumulators
/// are NOT serialized -- they rebuild naturally from continued training.
///
/// Requires the `serde-json` or `serde-bincode` feature.
#[cfg(any(feature = "serde-json", feature = "serde-bincode"))]
#[cfg_attr(
    docsrs,
    doc(cfg(any(feature = "serde-json", feature = "serde-bincode")))
)]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TreeSnapshot {
    /// Feature index chosen at each internal node (parallel to node arena).
    pub feature_idx: Vec<u32>,
    /// Split threshold at each internal node.
    pub threshold: Vec<f64>,
    /// Left child node index for each node (`u32::MAX` for leaves).
    pub left: Vec<u32>,
    /// Right child node index for each node (`u32::MAX` for leaves).
    pub right: Vec<u32>,
    /// Leaf prediction value for each node (only meaningful at leaves).
    pub leaf_value: Vec<f64>,
    /// Whether each node is a leaf.
    pub is_leaf: Vec<bool>,
    /// Depth of each node in the tree (root = 0).
    pub depth: Vec<u16>,
    /// Number of training samples that reached each node.
    pub sample_count: Vec<u64>,
    /// Total number of input features (`None` until first training sample).
    pub n_features: Option<usize>,
    /// Total samples seen by this tree since construction.
    pub samples_seen: u64,
    /// RNG state at the time of snapshot (for deterministic replay).
    pub rng_state: u64,
    /// Categorical split bitmasks. `None` entries are continuous splits.
    #[serde(default)]
    pub categorical_mask: Vec<Option<u64>>,
}

/// Serializable snapshot of a single boosting step.
///
/// Requires the `serde-json` or `serde-bincode` feature.
#[cfg(any(feature = "serde-json", feature = "serde-bincode"))]
#[cfg_attr(
    docsrs,
    doc(cfg(any(feature = "serde-json", feature = "serde-bincode")))
)]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StepSnapshot {
    /// Primary tree for this boosting step.
    pub tree: TreeSnapshot,
    /// Alternate tree being trained in parallel (present during drift warning period).
    pub alternate_tree: Option<TreeSnapshot>,
    /// Drift detector accumulated state (preserves warmup across save/load).
    #[serde(default)]
    pub drift_state: Option<crate::drift::state::DriftDetectorState>,
    /// Alternate drift detector state (if an alternate tree is training).
    #[serde(default)]
    pub alt_drift_state: Option<crate::drift::state::DriftDetectorState>,
}

/// Complete serializable state of an SGBT model.
///
/// Captures everything needed to reconstruct a trained model for prediction
/// and continued training. The loss function is stored as a [`LossType`] tag.
///
/// Requires the `serde-json` or `serde-bincode` feature.
#[cfg(any(feature = "serde-json", feature = "serde-bincode"))]
#[cfg_attr(
    docsrs,
    doc(cfg(any(feature = "serde-json", feature = "serde-bincode")))
)]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelState {
    /// Model configuration used to construct the ensemble.
    pub config: SGBTConfig,
    /// Loss function tag for runtime dispatch during restore.
    pub loss_type: LossType,
    /// Ensemble base prediction (warm-start bias).
    pub base_prediction: f64,
    /// Whether the base prediction has been initialized from data.
    pub base_initialized: bool,
    /// Targets buffered before the base prediction is fixed.
    pub initial_targets: Vec<f64>,
    /// Number of buffered targets used during initialization.
    pub initial_target_count: usize,
    /// Total training samples seen by the model.
    pub samples_seen: u64,
    /// RNG state at snapshot time (for deterministic replay).
    pub rng_state: u64,
    /// Per-step snapshots (one per boosting stage).
    pub steps: Vec<StepSnapshot>,
    /// Rolling mean absolute error for error-weighted sample importance.
    #[serde(default)]
    pub rolling_mean_error: f64,
    /// Per-step EWMA of contribution magnitude for quality pruning.
    #[serde(default)]
    pub contribution_ewma: Vec<f64>,
    /// Per-step consecutive low-contribution count for quality pruning.
    #[serde(default)]
    pub low_contrib_count: Vec<u64>,
    /// Rolling contribution sigma for adaptive_mts.
    #[serde(default)]
    pub rolling_contribution_sigma: f64,
}

/// Save an SGBT model to a JSON string.
///
/// Auto-detects the loss type from the model's loss function. For built-in
/// losses (Squared, Logistic, Huber, Softmax) this works automatically.
/// For custom losses, use [`save_model_with`] to supply the tag manually.
///
/// # Errors
///
/// Returns [`IrithyllError::Serialization`] if the loss type cannot be
/// auto-detected or if serialization fails.
#[cfg(feature = "serde-json")]
#[cfg_attr(docsrs, doc(cfg(feature = "serde-json")))]
pub fn save_model<L: Loss>(model: &SGBT<L>) -> Result<String> {
    let state = model.to_model_state()?;
    to_json_pretty(&state)
}

/// Save an SGBT model to a JSON string with an explicit loss type tag.
///
/// Use this for custom loss functions that don't implement `loss_type()`.
#[cfg(feature = "serde-json")]
#[cfg_attr(docsrs, doc(cfg(feature = "serde-json")))]
pub fn save_model_with<L: Loss>(model: &SGBT<L>, loss_type: LossType) -> Result<String> {
    let state = model.to_model_state_with(loss_type);
    to_json_pretty(&state)
}

/// Load an SGBT model from a JSON string.
///
/// Returns a [`DynSGBT`](crate::ensemble::DynSGBT) (`SGBT<Box<dyn Loss>>`)
/// because the concrete loss type is determined at runtime from the
/// serialized tag.
///
/// # Errors
///
/// Returns [`IrithyllError::Serialization`] if deserialization fails.
#[cfg(feature = "serde-json")]
#[cfg_attr(docsrs, doc(cfg(feature = "serde-json")))]
pub fn load_model(json: &str) -> Result<crate::ensemble::DynSGBT> {
    let state: ModelState = from_json(json)?;
    Ok(SGBT::from_model_state(state))
}

// ---------------------------------------------------------------------------
// MulticlassSGBT serialization
// ---------------------------------------------------------------------------

/// Serializable state for [`MulticlassSGBT`](crate::ensemble::multiclass::MulticlassSGBT).
///
/// Each class committee is stored as a full [`ModelState`] since each committee
/// is an independent `SGBT<SoftmaxLoss>`.
///
/// Requires the `serde-json` or `serde-bincode` feature.
#[cfg(any(feature = "serde-json", feature = "serde-bincode"))]
#[cfg_attr(
    docsrs,
    doc(cfg(any(feature = "serde-json", feature = "serde-bincode")))
)]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MulticlassModelState {
    /// Number of target classes.
    pub n_classes: usize,
    /// One [`ModelState`] per class committee.
    pub committees: Vec<ModelState>,
    /// Total training samples seen.
    pub samples_seen: u64,
}

/// Serialize a [`MulticlassModelState`] to JSON.
#[cfg(feature = "serde-json")]
#[cfg_attr(docsrs, doc(cfg(feature = "serde-json")))]
pub fn save_multiclass_model(state: &MulticlassModelState) -> Result<String> {
    to_json_pretty(state)
}

/// Deserialize a [`MulticlassModelState`] from JSON.
#[cfg(feature = "serde-json")]
#[cfg_attr(docsrs, doc(cfg(feature = "serde-json")))]
pub fn load_multiclass_model(json: &str) -> Result<MulticlassModelState> {
    from_json(json)
}

/// Serialize a [`MulticlassModelState`] to bincode bytes.
#[cfg(feature = "serde-bincode")]
#[cfg_attr(docsrs, doc(cfg(feature = "serde-bincode")))]
pub fn save_multiclass_model_bincode(state: &MulticlassModelState) -> Result<Vec<u8>> {
    to_bincode(state)
}

/// Deserialize a [`MulticlassModelState`] from bincode bytes.
#[cfg(feature = "serde-bincode")]
#[cfg_attr(docsrs, doc(cfg(feature = "serde-bincode")))]
pub fn load_multiclass_model_bincode(bytes: &[u8]) -> Result<MulticlassModelState> {
    from_bincode(bytes)
}

// ---------------------------------------------------------------------------
// BaggedSGBT serialization
// ---------------------------------------------------------------------------

/// Serializable state for [`BaggedSGBT`](crate::ensemble::bagged::BaggedSGBT).
///
/// Each bag is stored as a full [`ModelState`] since each bag is an
/// independent `SGBT<L>`.
///
/// Requires the `serde-json` or `serde-bincode` feature.
#[cfg(any(feature = "serde-json", feature = "serde-bincode"))]
#[cfg_attr(
    docsrs,
    doc(cfg(any(feature = "serde-json", feature = "serde-bincode")))
)]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BaggedModelState {
    /// Number of bootstrap bags.
    pub n_bags: usize,
    /// One [`ModelState`] per bag.
    pub bags: Vec<ModelState>,
    /// Total training samples seen.
    pub samples_seen: u64,
    /// RNG state at snapshot time.
    pub rng_state: u64,
    /// Original seed used to construct the bag sampler.
    pub seed: u64,
}

/// Serialize a [`BaggedModelState`] to JSON.
#[cfg(feature = "serde-json")]
#[cfg_attr(docsrs, doc(cfg(feature = "serde-json")))]
pub fn save_bagged_model(state: &BaggedModelState) -> Result<String> {
    to_json_pretty(state)
}

/// Deserialize a [`BaggedModelState`] from JSON.
#[cfg(feature = "serde-json")]
#[cfg_attr(docsrs, doc(cfg(feature = "serde-json")))]
pub fn load_bagged_model(json: &str) -> Result<BaggedModelState> {
    from_json(json)
}

/// Serialize a [`BaggedModelState`] to bincode bytes.
#[cfg(feature = "serde-bincode")]
#[cfg_attr(docsrs, doc(cfg(feature = "serde-bincode")))]
pub fn save_bagged_model_bincode(state: &BaggedModelState) -> Result<Vec<u8>> {
    to_bincode(state)
}

/// Deserialize a [`BaggedModelState`] from bincode bytes.
#[cfg(feature = "serde-bincode")]
#[cfg_attr(docsrs, doc(cfg(feature = "serde-bincode")))]
pub fn load_bagged_model_bincode(bytes: &[u8]) -> Result<BaggedModelState> {
    from_bincode(bytes)
}

// ---------------------------------------------------------------------------
// Bincode serialization (compact binary format)
// ---------------------------------------------------------------------------

/// Serialize a value to bincode bytes.
///
/// Requires the `serde-bincode` feature.
///
/// # Errors
///
/// Returns [`IrithyllError::Serialization`] if serialization fails.
#[cfg(feature = "serde-bincode")]
#[cfg_attr(docsrs, doc(cfg(feature = "serde-bincode")))]
pub fn to_bincode<T: serde::Serialize>(value: &T) -> Result<Vec<u8>> {
    bincode::serde::encode_to_vec(value, bincode::config::standard())
        .map_err(|e| IrithyllError::Serialization(e.to_string()))
}

/// Deserialize a value from bincode bytes.
///
/// Requires the `serde-bincode` feature.
///
/// # Errors
///
/// Returns [`IrithyllError::Serialization`] if deserialization fails.
#[cfg(feature = "serde-bincode")]
#[cfg_attr(docsrs, doc(cfg(feature = "serde-bincode")))]
pub fn from_bincode<T: serde::de::DeserializeOwned>(bytes: &[u8]) -> Result<T> {
    let (val, _) = bincode::serde::decode_from_slice(bytes, bincode::config::standard())
        .map_err(|e| IrithyllError::Serialization(e.to_string()))?;
    Ok(val)
}

/// Save an SGBT model to bincode bytes.
///
/// Compact binary format -- typically 3-5x smaller than JSON.
///
/// # Errors
///
/// Returns [`IrithyllError::Serialization`] if the loss type cannot be
/// auto-detected or if serialization fails.
#[cfg(feature = "serde-bincode")]
#[cfg_attr(docsrs, doc(cfg(feature = "serde-bincode")))]
pub fn save_model_bincode<L: Loss>(model: &SGBT<L>) -> Result<Vec<u8>> {
    let state = model.to_model_state()?;
    to_bincode(&state)
}

/// Load an SGBT model from bincode bytes.
///
/// Returns a [`DynSGBT`](crate::ensemble::DynSGBT) because the concrete
/// loss type is determined at runtime.
///
/// # Errors
///
/// Returns [`IrithyllError::Serialization`] if deserialization fails.
#[cfg(feature = "serde-bincode")]
#[cfg_attr(docsrs, doc(cfg(feature = "serde-bincode")))]
pub fn load_model_bincode(bytes: &[u8]) -> Result<crate::ensemble::DynSGBT> {
    let state: ModelState = from_bincode(bytes)?;
    Ok(SGBT::from_model_state(state))
}

// ---------------------------------------------------------------------------
// DistributionalSGBT serialization
// ---------------------------------------------------------------------------

/// Serializable state for [`DistributionalSGBT`](crate::ensemble::distributional::DistributionalSGBT).
///
/// Requires the `serde-json` or `serde-bincode` feature.
#[cfg(any(feature = "serde-json", feature = "serde-bincode"))]
#[cfg_attr(
    docsrs,
    doc(cfg(any(feature = "serde-json", feature = "serde-bincode")))
)]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DistributionalModelState {
    /// Model configuration shared by both location and scale ensembles.
    pub config: SGBTConfig,
    /// Boosting steps for the location (mean) ensemble.
    pub location_steps: Vec<StepSnapshot>,
    /// Boosting steps for the scale (uncertainty) ensemble.
    pub scale_steps: Vec<StepSnapshot>,
    /// Base prediction for the location ensemble.
    pub location_base: f64,
    /// Base prediction for the scale ensemble.
    pub scale_base: f64,
    /// Whether the base predictions have been initialized from data.
    pub base_initialized: bool,
    /// Targets buffered before base initialization is fixed.
    pub initial_targets: Vec<f64>,
    /// Number of buffered targets used during initialization.
    pub initial_target_count: usize,
    /// Total training samples seen.
    pub samples_seen: u64,
    /// RNG state at snapshot time.
    pub rng_state: u64,
    /// Whether uncertainty-modulated learning rate is active.
    pub uncertainty_modulated_lr: bool,
    /// Rolling mean of the scale (sigma) predictions.
    pub rolling_sigma_mean: f64,
    /// EWMA of squared prediction errors (for empirical σ mode).
    #[serde(default = "default_ewma_sq_err")]
    pub ewma_sq_err: f64,
    /// EWMA of honest_sigma (tree contribution std dev).
    #[serde(default)]
    pub rolling_honest_sigma_mean: f64,
}

fn default_ewma_sq_err() -> f64 {
    1.0
}

/// Serialize a [`DistributionalModelState`] to JSON.
#[cfg(feature = "serde-json")]
#[cfg_attr(docsrs, doc(cfg(feature = "serde-json")))]
pub fn save_distributional_model(state: &DistributionalModelState) -> Result<String> {
    to_json_pretty(state)
}

/// Deserialize a [`DistributionalModelState`] from JSON.
#[cfg(feature = "serde-json")]
#[cfg_attr(docsrs, doc(cfg(feature = "serde-json")))]
pub fn load_distributional_model(json: &str) -> Result<DistributionalModelState> {
    from_json(json)
}

/// Serialize a [`DistributionalModelState`] to bincode bytes.
#[cfg(feature = "serde-bincode")]
#[cfg_attr(docsrs, doc(cfg(feature = "serde-bincode")))]
pub fn save_distributional_model_bincode(state: &DistributionalModelState) -> Result<Vec<u8>> {
    to_bincode(state)
}

/// Deserialize a [`DistributionalModelState`] from bincode bytes.
#[cfg(feature = "serde-bincode")]
#[cfg_attr(docsrs, doc(cfg(feature = "serde-bincode")))]
pub fn load_distributional_model_bincode(bytes: &[u8]) -> Result<DistributionalModelState> {
    from_bincode(bytes)
}

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

    #[cfg(feature = "serde-json")]
    #[test]
    fn json_round_trip_sample() {
        let sample = Sample::new(vec![1.0, 2.0, 3.0], 4.0);
        let json = to_json(&sample).unwrap();
        let restored: Sample = from_json(&json).unwrap();
        assert_eq!(restored.features, sample.features);
        assert!((restored.target - sample.target).abs() < f64::EPSILON);
    }

    #[cfg(feature = "serde-json")]
    #[test]
    fn json_pretty_round_trip() {
        let sample = Sample::weighted(vec![1.0], 2.0, 0.5);
        let json = to_json_pretty(&sample).unwrap();
        assert!(json.contains('\n'));
        let restored: Sample = from_json(&json).unwrap();
        assert!((restored.weight - 0.5).abs() < f64::EPSILON);
    }

    #[cfg(feature = "serde-json")]
    #[test]
    fn json_bytes_round_trip() {
        let sample = Sample::new(vec![10.0, 20.0], 30.0);
        let bytes = to_json_bytes(&sample).unwrap();
        let restored: Sample = from_json_bytes(&bytes).unwrap();
        assert_eq!(restored.features, sample.features);
    }

    #[cfg(feature = "serde-json")]
    #[test]
    fn json_invalid_input_returns_error() {
        let result = from_json::<Sample>("not valid json");
        assert!(result.is_err());
        match result.unwrap_err() {
            IrithyllError::Serialization(msg) => {
                assert!(!msg.is_empty());
            }
            other => panic!("expected Serialization error, got {:?}", other),
        }
    }

    #[cfg(feature = "serde-json")]
    #[test]
    fn json_batch_samples() {
        let samples = vec![Sample::new(vec![1.0], 2.0), Sample::new(vec![3.0], 4.0)];
        let json = to_json(&samples).unwrap();
        let restored: Vec<Sample> = from_json(&json).unwrap();
        assert_eq!(restored.len(), 2);
    }

    #[cfg(feature = "serde-json")]
    #[test]
    fn multiclass_model_json_roundtrip() {
        use crate::ensemble::multiclass::MulticlassSGBT;
        use crate::SGBTConfig;

        let config = SGBTConfig::builder()
            .n_steps(5)
            .learning_rate(0.1)
            .grace_period(10)
            .max_depth(3)
            .initial_target_count(5)
            .build()
            .unwrap();

        let mut model = MulticlassSGBT::new(config, 3).unwrap();

        // Train on a simple 3-class problem
        for i in 0..60 {
            let x = i as f64 * 0.1;
            let class = (i % 3) as f64;
            model.train_one(&Sample::new(vec![x, x * 2.0], class));
        }

        // Serialize
        let state = model.to_multiclass_state();
        let json = save_multiclass_model(&state).unwrap();

        // Deserialize
        let loaded_state = load_multiclass_model(&json).unwrap();
        let restored = MulticlassSGBT::from_multiclass_state(loaded_state);

        // Verify predictions match
        let test_features = vec![vec![0.5, 1.0], vec![1.0, 2.0], vec![2.0, 4.0]];
        for features in &test_features {
            let orig_proba = model.predict_proba(features);
            let rest_proba = restored.predict_proba(features);
            assert_eq!(
                orig_proba.len(),
                rest_proba.len(),
                "probability vector lengths should match"
            );
            for (c, (o, r)) in orig_proba.iter().zip(rest_proba.iter()).enumerate() {
                assert!(
                    (o - r).abs() < 1e-10,
                    "multiclass JSON round-trip mismatch at class {}: {} vs {}",
                    c,
                    o,
                    r
                );
            }
        }

        // Verify metadata
        assert_eq!(model.n_classes(), restored.n_classes());
        assert_eq!(model.n_samples_seen(), restored.n_samples_seen());
    }

    #[cfg(feature = "serde-json")]
    #[test]
    fn bagged_model_json_roundtrip() {
        use crate::ensemble::bagged::BaggedSGBT;
        use crate::loss::squared::SquaredLoss;
        use crate::SGBTConfig;

        let config = SGBTConfig::builder()
            .n_steps(5)
            .learning_rate(0.1)
            .grace_period(10)
            .initial_target_count(5)
            .build()
            .unwrap();

        let mut model = BaggedSGBT::new(config, 3).unwrap();

        // Train on a simple regression problem
        for i in 0..100 {
            let x = i as f64 * 0.1;
            model.train_one(&Sample::new(vec![x], x * 2.0 + 1.0));
        }

        // Serialize
        let state = model.to_bagged_state().unwrap();
        let json = save_bagged_model(&state).unwrap();

        // Deserialize
        let loaded_state = load_bagged_model(&json).unwrap();
        let restored = BaggedSGBT::from_bagged_state(loaded_state, SquaredLoss);

        // Verify predictions match
        let test_points = [0.5, 1.0, 2.0, 3.0];
        for &x in &test_points {
            let orig = model.predict(&[x]);
            let rest = restored.predict(&[x]);
            assert!(
                (orig - rest).abs() < 1e-10,
                "bagged JSON round-trip mismatch at x={}: {} vs {}",
                x,
                orig,
                rest
            );
        }

        // Verify metadata
        assert_eq!(model.n_bags(), restored.n_bags());
        assert_eq!(model.n_samples_seen(), restored.n_samples_seen());
    }

    #[cfg(feature = "serde-json")]
    #[test]
    fn distributional_model_json_roundtrip() {
        use crate::ensemble::distributional::DistributionalSGBT;
        use crate::SGBTConfig;

        let config = SGBTConfig::builder()
            .n_steps(5)
            .learning_rate(0.1)
            .grace_period(10)
            .max_depth(3)
            .initial_target_count(10)
            .build()
            .unwrap();

        let mut model = DistributionalSGBT::new(config);

        // Train on a simple regression problem
        for i in 0..100 {
            let x = i as f64 * 0.1;
            model.train_one(&(vec![x], x.sin()));
        }

        // Serialize
        let state = model.to_distributional_state();
        let json = save_distributional_model(&state).unwrap();

        // Deserialize
        let loaded_state = load_distributional_model(&json).unwrap();
        let restored = DistributionalSGBT::from_distributional_state(loaded_state);

        // Verify predictions match
        let test_points = [0.5, 1.0, 2.0, 3.0];
        for &x in &test_points {
            let orig = model.predict(&[x]);
            let rest = restored.predict(&[x]);
            assert!(
                (orig.mu - rest.mu).abs() < 1e-10,
                "distributional JSON round-trip mu mismatch at x={}: {} vs {}",
                x,
                orig.mu,
                rest.mu
            );
            assert!(
                (orig.sigma - rest.sigma).abs() < 1e-10,
                "distributional JSON round-trip sigma mismatch at x={}: {} vs {}",
                x,
                orig.sigma,
                rest.sigma
            );
        }

        assert_eq!(model.n_samples_seen(), restored.n_samples_seen());
    }
}