rill_ml/models/

ftrl.rs

//! FTRL-Proximal online learning for sparse features.
//!
//! Implements the Follow-The-Regularized-Leader Proximal algorithm,
//! which is well-suited for high-dimensional sparse data. L1
//! regularization produces sparse weight vectors, and the per-coordinate
//! learning rate adapts to feature frequency.
//!
//! See: McMahan et al., "Ad Click Prediction: a View from the Trenches"
//! (KDD 2013).
//!
//! # Per-coordinate learning rate
//!
//! `eta_i = alpha / (beta + sqrt(n_i))`
//!
//! # Weight computation
//!
//! For feature `i`:
//!
//! ```text
//! if |z_i| <= lambda1:
//!     w_i = 0
//! else:
//!     w_i = -(z_i - sign(z_i) * lambda1) / (lambda2 + (beta + sqrt(n_i)) / alpha)
//! ```
//!
//! The intercept uses `lambda1 = 0` (no L1 regularization).
//!
//! # Failure atomicity
//!
//! [`FtrlRegressor::learn`] and [`FtrlClassifier::learn`] compute the next
//! state for every affected entry before mutating any field. If any
//! intermediate value is non-finite or the sample counter overflows, the
//! call returns `Err` and the model state is unchanged.
//!
//! # Dynamic feature growth
//!
//! By default the feature table grows without bound. Long-running
//! deployments handling untrusted streams should set
//! [`FtrlConfig::max_features`] to a finite value and pick a
//! [`NewFeaturePolicy`]. `None` preserves backwards compatibility but
//! allows memory growth proportional to the number of distinct
//! `FeatureId`s ever observed.

use crate::error::{RillError, checked_finite_add, checked_increment, ensure_finite};
use crate::loss::log_loss::sigmoid;
use crate::sparse::{FeatureId, SparseFeatures};
use crate::traits::{SparseClassifier, SparseRegressor};
use std::collections::BTreeMap;

/// Policy for handling new `FeatureId`s once [`FtrlConfig::max_features`]
/// has been reached.
///
/// `max_features` only constrains feature insertion; features already in
/// the model continue to train regardless of the policy.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum NewFeaturePolicy {
    /// Reject the entire `learn` call with [`RillError::InvalidState`] when
    /// adding the new features in the current sample would exceed
    /// `max_features`. The call is failure-atomic: no state changes.
    #[default]
    Reject,
    /// Keep training existing features but silently skip any new feature
    /// that would exceed `max_features`. The sample counter still
    /// advances and existing features still receive their gradient update.
    Ignore,
}

/// Configuration for FTRL models.
///
/// Controls the per-coordinate learning rate and regularization strengths.
/// All fields must be finite; `alpha` must be strictly positive and the
/// regularization parameters must be non-negative.
#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(serde::Serialize))]
pub struct FtrlConfig {
    /// Alpha: learning rate scaling. Must be `> 0`.
    pub alpha: f64,
    /// Beta: smoothing constant. Must be `>= 0`.
    pub beta: f64,
    /// L1 regularization strength. Must be `>= 0`.
    pub l1: f64,
    /// L2 regularization strength. Must be `>= 0`.
    pub l2: f64,
    /// Maximum number of distinct features the model will store.
    ///
    /// `None` (the default) allows unbounded growth, which is
    /// backwards-compatible but unsafe for long-running services that
    /// consume untrusted feature streams. Set to a finite value to
    /// bound memory and trigger the [`NewFeaturePolicy`].
    pub max_features: Option<usize>,
    /// Policy applied when `max_features` is set and a new `FeatureId`
    /// would exceed the cap.
    pub new_feature_policy: NewFeaturePolicy,
}

impl Default for FtrlConfig {
    fn default() -> Self {
        Self {
            alpha: 0.1,
            beta: 1.0,
            l1: 1.0,
            l2: 1.0,
            max_features: None,
            new_feature_policy: NewFeaturePolicy::default(),
        }
    }
}

impl FtrlConfig {
    /// Validate configuration parameters.
    fn validate(&self) -> Result<(), RillError> {
        ensure_finite("alpha", self.alpha)?;
        ensure_finite("beta", self.beta)?;
        ensure_finite("l1", self.l1)?;
        ensure_finite("l2", self.l2)?;
        if self.alpha <= 0.0 {
            return Err(RillError::InvalidParameter {
                name: "alpha",
                value: self.alpha,
            });
        }
        if self.beta < 0.0 {
            return Err(RillError::InvalidParameter {
                name: "beta",
                value: self.beta,
            });
        }
        if self.l1 < 0.0 {
            return Err(RillError::InvalidParameter {
                name: "l1",
                value: self.l1,
            });
        }
        if self.l2 < 0.0 {
            return Err(RillError::InvalidParameter {
                name: "l2",
                value: self.l2,
            });
        }
        if let Some(max_features) = self.max_features
            && max_features == 0
        {
            return Err(RillError::InvalidParameter {
                name: "max_features",
                value: 0.0,
            });
        }
        Ok(())
    }
}

#[cfg(feature = "serde")]
impl<'de> serde::Deserialize<'de> for FtrlConfig {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        #[derive(serde::Deserialize)]
        struct FtrlConfigState {
            alpha: f64,
            beta: f64,
            l1: f64,
            l2: f64,
            #[serde(default)]
            max_features: Option<usize>,
            #[serde(default)]
            new_feature_policy: NewFeaturePolicy,
        }

        let state = FtrlConfigState::deserialize(deserializer)?;
        let config = FtrlConfig {
            alpha: state.alpha,
            beta: state.beta,
            l1: state.l1,
            l2: state.l2,
            max_features: state.max_features,
            new_feature_policy: state.new_feature_policy,
        };
        config.validate().map_err(serde::de::Error::custom)?;
        Ok(config)
    }
}

/// Per-feature FTRL state.
///
/// Tracks the sum of (sigma-corrected) gradients `z` and the sum of squared
/// gradients `n`. The per-coordinate adaptive learning rate is derived from
/// `n`: features seen more frequently get smaller steps.
#[derive(Debug, Clone, Default)]
#[cfg_attr(feature = "serde", derive(serde::Serialize))]
pub struct FtrlParam {
    /// Sum of gradients (with sigma correction).
    z: f64,
    /// Sum of squared gradients. Must remain finite and non-negative.
    n: f64,
}

impl FtrlParam {
    /// Compute the FTRL weight given the config.
    ///
    /// Returns `0` when `|z| <= l1` (L1 soft-thresholding).
    fn weight(&self, config: &FtrlConfig) -> f64 {
        if self.z.abs() <= config.l1 {
            0.0
        } else {
            let sign = self.z.signum();
            let numerator = -(self.z - sign * config.l1);
            let denominator = config.l2 + (config.beta + self.n.sqrt()) / config.alpha;
            numerator / denominator
        }
    }

    /// Compute the intercept weight (no L1 regularization).
    ///
    /// Returns `0` when no gradient has been observed yet (`n == 0`),
    /// avoiding a potential `0 / 0` when `l2` and `beta` are both zero.
    fn intercept_weight(&self, config: &FtrlConfig) -> f64 {
        if self.n == 0.0 {
            0.0
        } else {
            let numerator = -self.z;
            let denominator = config.l2 + (config.beta + self.n.sqrt()) / config.alpha;
            numerator / denominator
        }
    }

    /// Compute the next `(z, n)` state after applying a gradient, without
    /// mutating `self`.
    ///
    /// `sigma = (sqrt(n_new) - sqrt(n_old)) / alpha`
    /// `z_new = z + g - sigma * w`
    /// `n_new = n + g^2`
    ///
    /// Every intermediate result is checked for finiteness so a partial
    /// overflow cannot poison state. The caller is responsible for
    /// committing the returned values atomically.
    fn next_updated(
        &self,
        gradient: f64,
        weight: f64,
        config: &FtrlConfig,
    ) -> Result<(f64, f64), RillError> {
        let gradient_sq = gradient * gradient;
        ensure_finite("ftrl_gradient_squared", gradient_sq)?;
        let n_new = checked_finite_add(self.n, gradient_sq, "ftrl_n_new")?;
        let sigma = (n_new.sqrt() - self.n.sqrt()) / config.alpha;
        ensure_finite("ftrl_sigma", sigma)?;
        let sigma_w = sigma * weight;
        ensure_finite("ftrl_sigma_weight", sigma_w)?;
        let z_delta = gradient - sigma_w;
        ensure_finite("ftrl_z_delta", z_delta)?;
        let z_new = checked_finite_add(self.z, z_delta, "ftrl_z_new")?;
        Ok((z_new, n_new))
    }

    /// Validate that `z` is finite and `n` is finite and non-negative.
    #[cfg_attr(not(feature = "serde"), allow(dead_code))]
    fn validate(&self) -> Result<(), RillError> {
        ensure_finite("ftrl_z", self.z)?;
        ensure_finite("ftrl_n", self.n)?;
        if self.n < 0.0 {
            return Err(RillError::InvalidState(format!(
                "ftrl n must be non-negative, got {0}",
                self.n
            )));
        }
        Ok(())
    }
}

#[cfg(feature = "serde")]
impl<'de> serde::Deserialize<'de> for FtrlParam {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        #[derive(serde::Deserialize)]
        struct FtrlParamState {
            z: f64,
            n: f64,
        }

        let state = FtrlParamState::deserialize(deserializer)?;
        let param = FtrlParam {
            z: state.z,
            n: state.n,
        };
        param.validate().map_err(serde::de::Error::custom)?;
        Ok(param)
    }
}

/// Compute the dot product `w · x` over sparse features.
///
/// Iterates only over the features present in `features` (not all stored
/// params), looking up each feature's current FTRL weight. Each
/// contribution and the running sum are checked for finiteness so a
/// single overflowing term cannot poison the prediction.
fn compute_dot(
    params: &BTreeMap<FeatureId, FtrlParam>,
    config: &FtrlConfig,
    features: &SparseFeatures,
) -> Result<f64, RillError> {
    if features.is_empty() {
        return Err(RillError::EmptyFeatures);
    }
    let mut dot = 0.0;
    for &(id, value) in features.values() {
        ensure_finite("sparse_value", value)?;
        if let Some(param) = params.get(&id) {
            let w = param.weight(config);
            ensure_finite("ftrl_weight", w)?;
            let contribution = w * value;
            ensure_finite("ftrl_dot_contribution", contribution)?;
            dot = checked_finite_add(dot, contribution, "ftrl_dot")?;
        }
    }
    Ok(dot)
}

/// FTRL regressor with squared loss.
///
/// Learns `y ≈ w · x + b` incrementally. The gradient of the squared loss
/// w.r.t. the prediction is `prediction - target`, so each feature's
/// gradient is `(prediction - target) * x_i`.
///
/// # Examples
///
/// ```
/// use rill_ml::models::{FtrlConfig, FtrlRegressor};
/// use rill_ml::sparse::SparseFeatures;
/// use rill_ml::SparseRegressor;
///
/// let mut model = FtrlRegressor::new(FtrlConfig::default()).unwrap();
/// let sf = SparseFeatures::from_sorted(vec![(0, 1.0), (1, 2.0)]).unwrap();
/// let _pred = model.predict(&sf).unwrap();
/// model.learn(&sf, 3.0).unwrap();
/// ```
#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(serde::Serialize))]
pub struct FtrlRegressor {
    config: FtrlConfig,
    params: BTreeMap<FeatureId, FtrlParam>,
    intercept: FtrlParam,
    samples_seen: u64,
}

impl FtrlRegressor {
    /// Create a new FTRL regressor.
    ///
    /// Returns an error if the configuration is invalid.
    pub fn new(config: FtrlConfig) -> Result<Self, RillError> {
        config.validate()?;
        Ok(Self {
            config,
            params: BTreeMap::new(),
            intercept: FtrlParam::default(),
            samples_seen: 0,
        })
    }

    /// The model configuration.
    pub const fn config(&self) -> &FtrlConfig {
        &self.config
    }

    /// Return the current non-zero feature weights, sorted by `FeatureId`.
    ///
    /// Features whose FTRL weight is exactly zero (due to L1
    /// soft-thresholding or never having been updated) are excluded.
    pub fn weights(&self) -> Vec<(FeatureId, f64)> {
        self.params
            .iter()
            .map(|(&id, param)| (id, param.weight(&self.config)))
            .filter(|&(_, w)| w != 0.0)
            .collect()
    }

    /// Compute the current intercept (bias) weight.
    pub fn intercept(&self) -> f64 {
        self.intercept.intercept_weight(&self.config)
    }

    /// Number of distinct features the model has seen.
    pub fn feature_count(&self) -> usize {
        self.params.len()
    }

    /// Compute the raw prediction `w · x + b` without updating state.
    fn predict_inner(&self, features: &SparseFeatures) -> Result<f64, RillError> {
        let dot = compute_dot(&self.params, &self.config, features)?;
        let intercept = self.intercept.intercept_weight(&self.config);
        ensure_finite("ftrl_intercept", intercept)?;
        Ok(dot + intercept)
    }
}

#[cfg(feature = "serde")]
impl<'de> serde::Deserialize<'de> for FtrlRegressor {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        #[derive(serde::Deserialize)]
        struct FtrlRegressorState {
            config: FtrlConfig,
            params: BTreeMap<FeatureId, FtrlParam>,
            intercept: FtrlParam,
            samples_seen: u64,
        }

        let state = FtrlRegressorState::deserialize(deserializer)?;
        let model = FtrlRegressor {
            config: state.config,
            params: state.params,
            intercept: state.intercept,
            samples_seen: state.samples_seen,
        };
        // config and each FtrlParam are already validated by their own
        // Deserialize impls; only the top-level invariants remain.
        model
            .validate_invariants()
            .map_err(serde::de::Error::custom)?;
        Ok(model)
    }
}

impl FtrlRegressor {
    #[cfg_attr(not(feature = "serde"), allow(dead_code))]
    fn validate_invariants(&self) -> Result<(), RillError> {
        // config and individual params are validated at deserialization;
        // nothing top-level to check beyond that.
        self.config.validate()?;
        self.intercept.validate()?;
        for param in self.params.values() {
            param.validate()?;
        }
        Ok(())
    }
}

impl SparseRegressor for FtrlRegressor {
    fn samples_seen(&self) -> u64 {
        self.samples_seen
    }

    fn predict(&self, features: &SparseFeatures) -> Result<f64, RillError> {
        self.predict_inner(features)
    }

    fn learn(&mut self, features: &SparseFeatures, target: f64) -> Result<(), RillError> {
        if features.is_empty() {
            return Err(RillError::EmptyFeatures);
        }
        ensure_finite("target", target)?;

        // Reserve the sample counter first. If it overflows, no state changes.
        let next_samples_seen = checked_increment(self.samples_seen, "samples_seen")?;

        let prediction = self.predict_inner(features)?;
        ensure_finite("ftrl_prediction", prediction)?;
        let grad = prediction - target;
        ensure_finite("ftrl_gradient", grad)?;

        // Pre-judge the max_features policy. We never insert a subset of
        // new features then fail: either all new features are accepted or
        // (under Ignore) all new features are skipped.
        let new_ids_count = features
            .values()
            .iter()
            .filter(|(id, _)| !self.params.contains_key(id))
            .count();
        let mut skip_new_features = false;
        if let Some(max_features) = self.config.max_features {
            let projected = self.params.len().saturating_add(new_ids_count);
            if projected > max_features {
                match self.config.new_feature_policy {
                    NewFeaturePolicy::Reject => {
                        return Err(RillError::InvalidState(format!(
                            "FTRL feature count {projected} exceeds max_features {max_features}"
                        )));
                    }
                    NewFeaturePolicy::Ignore => {
                        skip_new_features = true;
                    }
                }
            }
        }

        // Compute the next (z, n) for every feature without touching self.
        let mut updates: Vec<(FeatureId, f64, f64)> = Vec::with_capacity(features.len());
        for &(id, value) in features.values() {
            let g = grad * value;
            ensure_finite("ftrl_feature_gradient", g)?;

            let is_new = !self.params.contains_key(&id);
            if is_new && skip_new_features {
                continue;
            }

            let (new_z, new_n) = if let Some(param) = self.params.get(&id) {
                let w = param.weight(&self.config);
                param.next_updated(g, w, &self.config)?
            } else {
                let param = FtrlParam::default();
                let w = param.weight(&self.config);
                param.next_updated(g, w, &self.config)?
            };
            updates.push((id, new_z, new_n));
        }

        // Compute the next intercept state.
        let w_b = self.intercept.intercept_weight(&self.config);
        let (new_intercept_z, new_intercept_n) =
            self.intercept.next_updated(grad, w_b, &self.config)?;

        // Commit atomically. No failure path beyond this point.
        for (id, new_z, new_n) in updates {
            let param = self.params.entry(id).or_default();
            param.z = new_z;
            param.n = new_n;
        }
        self.intercept.z = new_intercept_z;
        self.intercept.n = new_intercept_n;
        self.samples_seen = next_samples_seen;

        Ok(())
    }

    fn reset(&mut self) {
        self.params.clear();
        self.intercept = FtrlParam::default();
        self.samples_seen = 0;
    }
}

/// FTRL binary classifier with log loss.
///
/// Predicts `P(y=1 | x) = sigmoid(w · x + b)`. The gradient of the log loss
/// w.r.t. the logit simplifies to `probability - target`, so each feature's
/// gradient is `(probability - target) * x_i`.
///
/// The returned probability lies in `[0, 1]`. Extreme logits can produce
/// exactly `0.0` or `1.0` after `sigmoid`; downstream consumers such as
/// [`crate::loss::log_loss::BinaryLogLoss`] clip internally.
///
/// # Examples
///
/// ```
/// use rill_ml::models::{FtrlClassifier, FtrlConfig};
/// use rill_ml::sparse::SparseFeatures;
/// use rill_ml::SparseClassifier;
///
/// let mut model = FtrlClassifier::new(FtrlConfig::default()).unwrap();
/// let sf = SparseFeatures::from_sorted(vec![(0, 1.0)]).unwrap();
/// let _proba = model.predict_proba(&sf).unwrap();
/// model.learn(&sf, true).unwrap();
/// ```
#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(serde::Serialize))]
pub struct FtrlClassifier {
    config: FtrlConfig,
    params: BTreeMap<FeatureId, FtrlParam>,
    intercept: FtrlParam,
    samples_seen: u64,
}

impl FtrlClassifier {
    /// Create a new FTRL classifier.
    ///
    /// Returns an error if the configuration is invalid.
    pub fn new(config: FtrlConfig) -> Result<Self, RillError> {
        config.validate()?;
        Ok(Self {
            config,
            params: BTreeMap::new(),
            intercept: FtrlParam::default(),
            samples_seen: 0,
        })
    }

    /// The model configuration.
    pub const fn config(&self) -> &FtrlConfig {
        &self.config
    }

    /// Return the current non-zero feature weights, sorted by `FeatureId`.
    ///
    /// Features whose FTRL weight is exactly zero (due to L1
    /// soft-thresholding or never having been updated) are excluded.
    pub fn weights(&self) -> Vec<(FeatureId, f64)> {
        self.params
            .iter()
            .map(|(&id, param)| (id, param.weight(&self.config)))
            .filter(|&(_, w)| w != 0.0)
            .collect()
    }

    /// Compute the current intercept (bias) weight.
    pub fn intercept(&self) -> f64 {
        self.intercept.intercept_weight(&self.config)
    }

    /// Number of distinct features the model has seen.
    pub fn feature_count(&self) -> usize {
        self.params.len()
    }

    /// Compute the probability `sigmoid(w · x + b)` without updating state.
    fn predict_proba_inner(&self, features: &SparseFeatures) -> Result<f64, RillError> {
        let dot = compute_dot(&self.params, &self.config, features)?;
        let intercept = self.intercept.intercept_weight(&self.config);
        ensure_finite("ftrl_intercept", intercept)?;
        let logit = dot + intercept;
        ensure_finite("ftrl_logit", logit)?;
        Ok(sigmoid(logit))
    }
}

#[cfg(feature = "serde")]
impl<'de> serde::Deserialize<'de> for FtrlClassifier {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        #[derive(serde::Deserialize)]
        struct FtrlClassifierState {
            config: FtrlConfig,
            params: BTreeMap<FeatureId, FtrlParam>,
            intercept: FtrlParam,
            samples_seen: u64,
        }

        let state = FtrlClassifierState::deserialize(deserializer)?;
        let model = FtrlClassifier {
            config: state.config,
            params: state.params,
            intercept: state.intercept,
            samples_seen: state.samples_seen,
        };
        model
            .validate_invariants()
            .map_err(serde::de::Error::custom)?;
        Ok(model)
    }
}

impl FtrlClassifier {
    #[cfg_attr(not(feature = "serde"), allow(dead_code))]
    fn validate_invariants(&self) -> Result<(), RillError> {
        self.config.validate()?;
        self.intercept.validate()?;
        for param in self.params.values() {
            param.validate()?;
        }
        Ok(())
    }
}

impl SparseClassifier for FtrlClassifier {
    fn samples_seen(&self) -> u64 {
        self.samples_seen
    }

    fn predict_proba(&self, features: &SparseFeatures) -> Result<f64, RillError> {
        self.predict_proba_inner(features)
    }

    fn learn(&mut self, features: &SparseFeatures, target: bool) -> Result<(), RillError> {
        if features.is_empty() {
            return Err(RillError::EmptyFeatures);
        }

        let next_samples_seen = checked_increment(self.samples_seen, "samples_seen")?;

        let probability = self.predict_proba_inner(features)?;
        ensure_finite("ftrl_probability", probability)?;
        let y = if target { 1.0 } else { 0.0 };
        let grad = probability - y;
        ensure_finite("ftrl_gradient", grad)?;

        let new_ids_count = features
            .values()
            .iter()
            .filter(|(id, _)| !self.params.contains_key(id))
            .count();
        let mut skip_new_features = false;
        if let Some(max_features) = self.config.max_features {
            let projected = self.params.len().saturating_add(new_ids_count);
            if projected > max_features {
                match self.config.new_feature_policy {
                    NewFeaturePolicy::Reject => {
                        return Err(RillError::InvalidState(format!(
                            "FTRL feature count {projected} exceeds max_features {max_features}"
                        )));
                    }
                    NewFeaturePolicy::Ignore => {
                        skip_new_features = true;
                    }
                }
            }
        }

        let mut updates: Vec<(FeatureId, f64, f64)> = Vec::with_capacity(features.len());
        for &(id, value) in features.values() {
            let g = grad * value;
            ensure_finite("ftrl_feature_gradient", g)?;

            let is_new = !self.params.contains_key(&id);
            if is_new && skip_new_features {
                continue;
            }

            let (new_z, new_n) = if let Some(param) = self.params.get(&id) {
                let w = param.weight(&self.config);
                param.next_updated(g, w, &self.config)?
            } else {
                let param = FtrlParam::default();
                let w = param.weight(&self.config);
                param.next_updated(g, w, &self.config)?
            };
            updates.push((id, new_z, new_n));
        }

        let w_b = self.intercept.intercept_weight(&self.config);
        let (new_intercept_z, new_intercept_n) =
            self.intercept.next_updated(grad, w_b, &self.config)?;

        for (id, new_z, new_n) in updates {
            let param = self.params.entry(id).or_default();
            param.z = new_z;
            param.n = new_n;
        }
        self.intercept.z = new_intercept_z;
        self.intercept.n = new_intercept_n;
        self.samples_seen = next_samples_seen;

        Ok(())
    }

    fn reset(&mut self) {
        self.params.clear();
        self.intercept = FtrlParam::default();
        self.samples_seen = 0;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rand::SeedableRng;

    // -----------------------------------------------------------------
    // FtrlRegressor tests
    // -----------------------------------------------------------------

    #[test]
    fn cold_start_returns_zero() {
        let model = FtrlRegressor::new(FtrlConfig::default()).unwrap();
        let sf = SparseFeatures::from_sorted(vec![(0, 1.0)]).unwrap();
        let pred = model.predict(&sf).unwrap();
        assert!(pred.abs() < 1e-12);
    }

    #[test]
    fn learn_linear_data_converges() {
        // y = 2 * x, single feature
        let mut model = FtrlRegressor::new(FtrlConfig {
            alpha: 0.5,
            beta: 1.0,
            l1: 0.0,
            l2: 0.0,
            max_features: None,
            new_feature_policy: NewFeaturePolicy::default(),
        })
        .unwrap();
        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(42);
        let mut first_err = 0.0;
        let mut last_err = 0.0;
        for i in 0..500 {
            let x = rand::Rng::gen_range(&mut rng, -1.0..1.0);
            let y = 2.0 * x;
            let sf = SparseFeatures::from_sorted(vec![(0, x)]).unwrap();
            let pred = model.predict(&sf).unwrap();
            let err = (pred - y).abs();
            if i < 10 {
                first_err += err;
            }
            if i >= 490 {
                last_err += err;
            }
            model.learn(&sf, y).unwrap();
        }
        assert!(last_err < first_err, "error should decrease");
        let weights = model.weights();
        assert_eq!(weights.len(), 1);
        assert!(
            (weights[0].1 - 2.0).abs() < 0.5,
            "weight should approach 2.0"
        );
    }

    #[test]
    fn l1_produces_sparse_weights() {
        // High L1 should drive most weights to zero.
        let mut model = FtrlRegressor::new(FtrlConfig {
            alpha: 0.1,
            beta: 1.0,
            l1: 100.0,
            l2: 0.0,
            max_features: None,
            new_feature_policy: NewFeaturePolicy::default(),
        })
        .unwrap();
        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(1);
        for _ in 0..200 {
            let x1 = rand::Rng::gen_range(&mut rng, -1.0..1.0);
            let x2 = rand::Rng::gen_range(&mut rng, -1.0..1.0);
            let y = 0.5 * x1;
            let sf = SparseFeatures::from_sorted(vec![(0, x1), (1, x2)]).unwrap();
            model.learn(&sf, y).unwrap();
        }
        let weights = model.weights();
        // With very high L1, all weights should be zero.
        assert!(
            weights.is_empty(),
            "weights should all be zero, got {weights:?}"
        );
    }

    #[test]
    fn dynamic_features() {
        let mut model = FtrlRegressor::new(FtrlConfig::default()).unwrap();
        assert_eq!(model.feature_count(), 0);
        let sf1 = SparseFeatures::from_sorted(vec![(0, 1.0)]).unwrap();
        model.learn(&sf1, 1.0).unwrap();
        assert_eq!(model.feature_count(), 1);
        // A new feature id appears.
        let sf2 = SparseFeatures::from_sorted(vec![(5, 2.0)]).unwrap();
        model.learn(&sf2, 2.0).unwrap();
        assert_eq!(model.feature_count(), 2);
        // Feature 0 still present.
        assert!(model.params.contains_key(&0));
        assert!(model.params.contains_key(&5));
    }

    #[test]
    fn predict_does_not_update_state() {
        let mut model = FtrlRegressor::new(FtrlConfig::default()).unwrap();
        let sf = SparseFeatures::from_sorted(vec![(0, 1.0)]).unwrap();
        let _ = model.predict(&sf).unwrap();
        assert_eq!(model.samples_seen(), 0);
        assert_eq!(model.feature_count(), 0);
        // Learn once, then predict again.
        model.learn(&sf, 1.0).unwrap();
        let count_after_learn = model.feature_count();
        let _ = model.predict(&sf).unwrap();
        assert_eq!(model.feature_count(), count_after_learn);
        assert_eq!(model.samples_seen(), 1);
    }

    #[test]
    fn non_finite_value_rejected() {
        let model = FtrlRegressor::new(FtrlConfig::default()).unwrap();
        // SparseFeatures::from_sorted rejects non-finite values at construction.
        assert!(SparseFeatures::from_sorted(vec![(0, f64::NAN)]).is_err());
        assert!(SparseFeatures::from_sorted(vec![(0, f64::INFINITY)]).is_err());
        assert!(SparseFeatures::from_sorted(vec![(0, f64::NEG_INFINITY)]).is_err());
        let sf = SparseFeatures::from_sorted(vec![(0, 1.0)]).unwrap();
        assert!(model.predict(&sf).is_ok());
    }

    #[test]
    fn non_finite_target_rejected() {
        let mut model = FtrlRegressor::new(FtrlConfig::default()).unwrap();
        let sf = SparseFeatures::from_sorted(vec![(0, 1.0)]).unwrap();
        assert!(model.learn(&sf, f64::NAN).is_err());
        assert!(model.learn(&sf, f64::INFINITY).is_err());
        assert!(model.learn(&sf, f64::NEG_INFINITY).is_err());
        // State should not change on error.
        assert_eq!(model.samples_seen(), 0);
    }

    #[test]
    fn empty_features_rejected() {
        let mut model = FtrlRegressor::new(FtrlConfig::default()).unwrap();
        let sf = SparseFeatures::new();
        assert!(model.predict(&sf).is_err());
        assert!(model.learn(&sf, 1.0).is_err());
    }

    #[test]
    fn reset_clears_state() {
        let mut model = FtrlRegressor::new(FtrlConfig::default()).unwrap();
        let sf = SparseFeatures::from_sorted(vec![(0, 1.0), (1, 2.0)]).unwrap();
        model.learn(&sf, 3.0).unwrap();
        model.learn(&sf, 3.0).unwrap();
        assert_eq!(model.samples_seen(), 2);
        assert_eq!(model.feature_count(), 2);
        model.reset();
        assert_eq!(model.samples_seen(), 0);
        assert_eq!(model.feature_count(), 0);
        assert!(model.predict(&sf).unwrap().abs() < 1e-12);
    }

    #[test]
    fn invalid_config_rejected() {
        assert!(
            FtrlRegressor::new(FtrlConfig {
                alpha: 0.0,
                ..FtrlConfig::default()
            })
            .is_err()
        );
        assert!(
            FtrlRegressor::new(FtrlConfig {
                alpha: -1.0,
                ..FtrlConfig::default()
            })
            .is_err()
        );
        assert!(
            FtrlRegressor::new(FtrlConfig {
                beta: -1.0,
                ..FtrlConfig::default()
            })
            .is_err()
        );
        assert!(
            FtrlRegressor::new(FtrlConfig {
                l1: -1.0,
                ..FtrlConfig::default()
            })
            .is_err()
        );
        assert!(
            FtrlRegressor::new(FtrlConfig {
                l2: -1.0,
                ..FtrlConfig::default()
            })
            .is_err()
        );
        assert!(
            FtrlRegressor::new(FtrlConfig {
                alpha: f64::NAN,
                ..FtrlConfig::default()
            })
            .is_err()
        );
        assert!(
            FtrlRegressor::new(FtrlConfig {
                max_features: Some(0),
                ..FtrlConfig::default()
            })
            .is_err()
        );
    }

    #[test]
    #[cfg(feature = "serde")]
    fn serde_roundtrip() {
        let mut model = FtrlRegressor::new(FtrlConfig {
            alpha: 0.2,
            beta: 0.5,
            l1: 0.5,
            l2: 0.5,
            max_features: Some(100),
            new_feature_policy: NewFeaturePolicy::Reject,
        })
        .unwrap();
        let sf = SparseFeatures::from_sorted(vec![(0, 1.0), (3, 2.0)]).unwrap();
        model.learn(&sf, 5.0).unwrap();
        let json = serde_json::to_string(&model).unwrap();
        let restored: FtrlRegressor = serde_json::from_str(&json).unwrap();
        assert_eq!(restored.samples_seen(), model.samples_seen());
        assert_eq!(restored.feature_count(), model.feature_count());
        let pred_orig = model.predict(&sf).unwrap();
        let pred_restored = restored.predict(&sf).unwrap();
        assert!((pred_orig - pred_restored).abs() < 1e-12);
    }

    #[test]
    fn weights_returns_nonzero_only() {
        let mut model = FtrlRegressor::new(FtrlConfig {
            alpha: 0.5,
            beta: 1.0,
            l1: 0.0,
            l2: 0.0,
            max_features: None,
            new_feature_policy: NewFeaturePolicy::default(),
        })
        .unwrap();
        // Learn feature 0 strongly, feature 1 barely.
        let sf = SparseFeatures::from_sorted(vec![(0, 1.0), (1, 0.0001)]).unwrap();
        for _ in 0..50 {
            model.learn(&sf, 1.0).unwrap();
        }
        let weights = model.weights();
        // All returned weights should be non-zero.
        for &(_, w) in &weights {
            assert!(w != 0.0);
        }
        // Feature 0 should be in the list.
        assert!(weights.iter().any(|&(id, _)| id == 0));
    }

    #[test]
    fn multiple_features() {
        // y = 1.0 * x0 + (-1.0) * x1 + 0.5
        let mut model = FtrlRegressor::new(FtrlConfig {
            alpha: 0.5,
            beta: 1.0,
            l1: 0.0,
            l2: 0.0,
            max_features: None,
            new_feature_policy: NewFeaturePolicy::default(),
        })
        .unwrap();
        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(99);
        for _ in 0..500 {
            let x0 = rand::Rng::gen_range(&mut rng, -1.0..1.0);
            let x1 = rand::Rng::gen_range(&mut rng, -1.0..1.0);
            let y = 1.0 * x0 - 1.0 * x1 + 0.5;
            let sf = SparseFeatures::from_sorted(vec![(0, x0), (1, x1)]).unwrap();
            model.learn(&sf, y).unwrap();
        }
        let weights = model.weights();
        assert_eq!(weights.len(), 2);
        let w0 = weights
            .iter()
            .find(|&&(id, _)| id == 0)
            .map(|&(_, w)| w)
            .unwrap();
        let w1 = weights
            .iter()
            .find(|&&(id, _)| id == 1)
            .map(|&(_, w)| w)
            .unwrap();
        assert!((w0 - 1.0).abs() < 0.5, "w0 should approach 1.0, got {w0}");
        assert!((w1 + 1.0).abs() < 0.5, "w1 should approach -1.0, got {w1}");
        assert!(
            (model.intercept() - 0.5).abs() < 0.5,
            "intercept should approach 0.5"
        );
    }

    #[test]
    fn intercept_learned() {
        // y = 3.0 (constant), single feature with value 0.0 so that only
        // the intercept can learn (feature gradient is always 0).
        let mut model = FtrlRegressor::new(FtrlConfig {
            alpha: 0.5,
            beta: 1.0,
            l1: 0.0,
            l2: 0.0,
            max_features: None,
            new_feature_policy: NewFeaturePolicy::default(),
        })
        .unwrap();
        let sf = SparseFeatures::from_sorted(vec![(0, 0.0)]).unwrap();
        for _ in 0..300 {
            model.learn(&sf, 3.0).unwrap();
        }
        let pred = model.predict(&sf).unwrap();
        assert!(
            (pred - 3.0).abs() < 0.5,
            "prediction should approach 3.0, got {pred}"
        );
        assert!(
            (model.intercept() - 3.0).abs() < 0.5,
            "intercept should approach 3.0"
        );
        // Feature weight should be 0 (never updated since x=0).
        assert!(model.weights().is_empty());
    }

    #[test]
    fn high_dim_sparse() {
        // 1000 possible features, only 5 active per sample.
        // Target is a linear combination of the active features.
        let mut model = FtrlRegressor::new(FtrlConfig {
            alpha: 0.3,
            beta: 1.0,
            l1: 0.0,
            l2: 0.0,
            max_features: None,
            new_feature_policy: NewFeaturePolicy::default(),
        })
        .unwrap();
        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(7);
        // True weights for features 0..5.
        let true_w = [1.0, -0.5, 2.0, 0.3, -1.5];
        let mut first_err = 0.0;
        let mut last_err = 0.0;
        for i in 0..2000 {
            let mut active: Vec<(FeatureId, f64)> = Vec::with_capacity(5);
            for (j, &w) in true_w.iter().enumerate() {
                let x = rand::Rng::gen_range(&mut rng, -1.0..1.0);
                active.push((j as u64, x * w));
            }
            // Add some noise features with zero contribution.
            for k in 5..10 {
                let x = rand::Rng::gen_range(&mut rng, -1.0..1.0);
                active.push((k as u64 + 100, x));
            }
            active.sort_by_key(|(id, _)| *id);
            let sf = SparseFeatures::from_sorted(active.clone()).unwrap();
            let y: f64 = active.iter().take(5).map(|(_, v)| v).sum();
            let pred = model.predict(&sf).unwrap();
            let err = (pred - y).abs();
            if i < 20 {
                first_err += err;
            }
            if i >= 1980 {
                last_err += err;
            }
            model.learn(&sf, y).unwrap();
        }
        assert!(
            last_err < first_err,
            "error should decrease in high-dim sparse"
        );
    }

    // -----------------------------------------------------------------
    // FtrlRegressor: failure atomicity and overflow (ML-001/002)
    // -----------------------------------------------------------------

    #[test]
    fn regressor_overflow_does_not_mutate_state() {
        // Finite inputs but intermediate `gradient * value` or
        // `gradient^2` overflows. The whole learn call must fail and
        // leave the model untouched.
        let mut model = FtrlRegressor::new(FtrlConfig {
            alpha: 0.1,
            beta: 1.0,
            l1: 0.0,
            l2: 0.0,
            max_features: None,
            new_feature_policy: NewFeaturePolicy::default(),
        })
        .unwrap();
        let sf = SparseFeatures::from_sorted(vec![(0, 1.0), (1, 1e200)]).unwrap();
        let result = model.learn(&sf, 1e100);
        assert!(result.is_err(), "expected overflow error, got {result:?}");
        assert_eq!(model.samples_seen(), 0);
        assert_eq!(model.feature_count(), 0);
        assert!(model.params.is_empty());
        assert_eq!(model.intercept.z, 0.0);
        assert_eq!(model.intercept.n, 0.0);
    }

    #[test]
    fn regressor_partial_update_is_atomic() {
        // Two features in one sample. Feature 0 would succeed on its own,
        // feature 1 overflows. Neither may be committed.
        let mut model = FtrlRegressor::new(FtrlConfig {
            alpha: 0.1,
            beta: 1.0,
            l1: 0.0,
            l2: 0.0,
            max_features: None,
            new_feature_policy: NewFeaturePolicy::default(),
        })
        .unwrap();
        let sf = SparseFeatures::from_sorted(vec![(0, 1.0), (1, 1e200)]).unwrap();
        assert!(model.learn(&sf, 1e100).is_err());
        // Feature 0 must NOT be inserted.
        assert!(!model.params.contains_key(&0));
        assert!(!model.params.contains_key(&1));
        assert_eq!(model.samples_seen(), 0);
    }

    #[test]
    #[cfg(feature = "serde")]
    fn regressor_samples_seen_overflow_is_atomic() {
        let json = format!(
            "{{\"config\":{{\"alpha\":0.1,\"beta\":1.0,\"l1\":1.0,\"l2\":1.0,\"max_features\":null,\"new_feature_policy\":\"Reject\"}},\"params\":{{}},\"intercept\":{{\"z\":0.0,\"n\":0.0}},\"samples_seen\":{}}}",
            u64::MAX
        );
        let mut model: FtrlRegressor = serde_json::from_str(&json).unwrap();
        let sf = SparseFeatures::from_sorted(vec![(0, 1.0)]).unwrap();
        let result = model.learn(&sf, 1.0);
        assert!(result.is_err(), "expected counter overflow");
        assert_eq!(model.samples_seen(), u64::MAX);
        assert_eq!(model.feature_count(), 0);
        assert_eq!(model.intercept.z, 0.0);
        assert_eq!(model.intercept.n, 0.0);
    }

    // -----------------------------------------------------------------
    // FtrlRegressor: max_features boundary (ML-003)
    // -----------------------------------------------------------------

    #[test]
    fn regressor_max_features_reject_at_limit() {
        let mut model = FtrlRegressor::new(FtrlConfig {
            alpha: 0.5,
            beta: 1.0,
            l1: 0.0,
            l2: 0.0,
            max_features: Some(2),
            new_feature_policy: NewFeaturePolicy::Reject,
        })
        .unwrap();
        // Reach exactly the limit.
        let sf = SparseFeatures::from_sorted(vec![(0, 1.0), (1, 2.0)]).unwrap();
        model.learn(&sf, 1.0).unwrap();
        assert_eq!(model.feature_count(), 2);
        // One more new feature: Reject must fail atomically.
        let sf_new = SparseFeatures::from_sorted(vec![(2, 3.0)]).unwrap();
        assert!(model.learn(&sf_new, 1.0).is_err());
        assert_eq!(model.feature_count(), 2);
        assert_eq!(model.samples_seen(), 1);
        // Existing features still train.
        let sf_existing = SparseFeatures::from_sorted(vec![(0, 1.0), (1, 2.0)]).unwrap();
        model.learn(&sf_existing, 1.0).unwrap();
        assert_eq!(model.feature_count(), 2);
        assert_eq!(model.samples_seen(), 2);
    }

    #[test]
    fn regressor_max_features_ignore_skips_new() {
        let mut model = FtrlRegressor::new(FtrlConfig {
            alpha: 0.5,
            beta: 1.0,
            l1: 0.0,
            l2: 0.0,
            max_features: Some(2),
            new_feature_policy: NewFeaturePolicy::Ignore,
        })
        .unwrap();
        let sf = SparseFeatures::from_sorted(vec![(0, 1.0), (1, 2.0)]).unwrap();
        model.learn(&sf, 1.0).unwrap();
        // Sample with one existing and one new feature. Under Ignore the
        // new feature is skipped, the existing one is updated, and the
        // counter still advances.
        let sf_mixed = SparseFeatures::from_sorted(vec![(0, 1.0), (2, 3.0)]).unwrap();
        model.learn(&sf_mixed, 1.0).unwrap();
        assert_eq!(model.feature_count(), 2);
        assert!(!model.params.contains_key(&2));
        assert_eq!(model.samples_seen(), 2);
    }

    #[test]
    fn regressor_max_features_multi_new_prejudge() {
        let mut model = FtrlRegressor::new(FtrlConfig {
            alpha: 0.5,
            beta: 1.0,
            l1: 0.0,
            l2: 0.0,
            max_features: Some(2),
            new_feature_policy: NewFeaturePolicy::Reject,
        })
        .unwrap();
        // A single sample with three new features. Reject fails atomically
        // without inserting any subset.
        let sf = SparseFeatures::from_sorted(vec![(0, 1.0), (1, 2.0), (2, 3.0)]).unwrap();
        assert!(model.learn(&sf, 1.0).is_err());
        assert_eq!(model.feature_count(), 0);
        assert_eq!(model.samples_seen(), 0);
    }

    // -----------------------------------------------------------------
    // FtrlRegressor: serde validation (ML-004)
    // -----------------------------------------------------------------

    #[test]
    #[cfg(feature = "serde")]
    fn regressor_serde_rejects_negative_n() {
        let json = "{\"config\":{\"alpha\":0.1,\"beta\":1.0,\"l1\":1.0,\"l2\":1.0,\"max_features\":null,\"new_feature_policy\":\"Reject\"},\"params\":{\"0\":{\"z\":1.0,\"n\":-1.0}},\"intercept\":{\"z\":0.0,\"n\":0.0},\"samples_seen\":0}";
        let result: Result<FtrlRegressor, _> = serde_json::from_str(json);
        assert!(result.is_err(), "negative n must be rejected");
    }

    #[test]
    #[cfg(feature = "serde")]
    fn regressor_serde_rejects_invalid_config() {
        let json = "{\"config\":{\"alpha\":-1.0,\"beta\":1.0,\"l1\":1.0,\"l2\":1.0,\"max_features\":null,\"new_feature_policy\":\"Reject\"},\"params\":{},\"intercept\":{\"z\":0.0,\"n\":0.0},\"samples_seen\":0}";
        let result: Result<FtrlRegressor, _> = serde_json::from_str(json);
        assert!(result.is_err(), "invalid alpha must be rejected");
    }

    #[test]
    #[cfg(feature = "serde")]
    fn regressor_serde_accepts_missing_optional_fields() {
        // Old state without max_features/new_feature_policy must still load.
        let json = "{\"config\":{\"alpha\":0.1,\"beta\":1.0,\"l1\":1.0,\"l2\":1.0},\"params\":{},\"intercept\":{\"z\":0.0,\"n\":0.0},\"samples_seen\":0}";
        let model: FtrlRegressor = serde_json::from_str(json).unwrap();
        assert!(model.config().max_features.is_none());
        assert_eq!(model.config().new_feature_policy, NewFeaturePolicy::Reject);
    }

    // -----------------------------------------------------------------
    // FtrlClassifier tests
    // -----------------------------------------------------------------

    #[test]
    fn cold_start_returns_0_5() {
        let model = FtrlClassifier::new(FtrlConfig::default()).unwrap();
        let sf = SparseFeatures::from_sorted(vec![(0, 1.0)]).unwrap();
        let p = model.predict_proba(&sf).unwrap();
        assert!((p - 0.5).abs() < 1e-12, "cold start should predict 0.5");
    }

    #[test]
    fn learn_separable_data() {
        // Linearly separable: class 1 when x0 > 0.
        let mut model = FtrlClassifier::new(FtrlConfig {
            alpha: 0.5,
            beta: 1.0,
            l1: 0.0,
            l2: 0.0,
            max_features: None,
            new_feature_policy: NewFeaturePolicy::default(),
        })
        .unwrap();
        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(3);
        for _ in 0..1000 {
            let x0 = rand::Rng::gen_range(&mut rng, -2.0..2.0);
            let x1 = rand::Rng::gen_range(&mut rng, -1.0..1.0);
            let y = x0 > 0.0;
            let sf = SparseFeatures::from_sorted(vec![(0, x0), (1, x1)]).unwrap();
            model.learn(&sf, y).unwrap();
        }
        let p_pos = model
            .predict_proba(&SparseFeatures::from_sorted(vec![(0, 2.0), (1, 0.0)]).unwrap())
            .unwrap();
        let p_neg = model
            .predict_proba(&SparseFeatures::from_sorted(vec![(0, -2.0), (1, 0.0)]).unwrap())
            .unwrap();
        assert!(p_pos > 0.7, "p_pos should be high, got {p_pos}");
        assert!(p_neg < 0.3, "p_neg should be low, got {p_neg}");
    }

    #[test]
    fn classifier_l1_produces_sparse_weights() {
        let mut model = FtrlClassifier::new(FtrlConfig {
            alpha: 0.1,
            beta: 1.0,
            l1: 100.0,
            l2: 0.0,
            max_features: None,
            new_feature_policy: NewFeaturePolicy::default(),
        })
        .unwrap();
        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(5);
        for _ in 0..200 {
            let x0 = rand::Rng::gen_range(&mut rng, -1.0..1.0);
            let x1 = rand::Rng::gen_range(&mut rng, -1.0..1.0);
            let y = x0 > 0.0;
            let sf = SparseFeatures::from_sorted(vec![(0, x0), (1, x1)]).unwrap();
            model.learn(&sf, y).unwrap();
        }
        let weights = model.weights();
        assert!(
            weights.is_empty(),
            "weights should all be zero with high L1, got {weights:?}"
        );
    }

    #[test]
    fn classifier_dynamic_features() {
        let mut model = FtrlClassifier::new(FtrlConfig::default()).unwrap();
        assert_eq!(model.feature_count(), 0);
        let sf1 = SparseFeatures::from_sorted(vec![(0, 1.0)]).unwrap();
        model.learn(&sf1, true).unwrap();
        assert_eq!(model.feature_count(), 1);
        let sf2 = SparseFeatures::from_sorted(vec![(10, 1.0)]).unwrap();
        model.learn(&sf2, false).unwrap();
        assert_eq!(model.feature_count(), 2);
    }

    #[test]
    fn classifier_predict_does_not_update_state() {
        let mut model = FtrlClassifier::new(FtrlConfig::default()).unwrap();
        let sf = SparseFeatures::from_sorted(vec![(0, 1.0)]).unwrap();
        let _ = model.predict_proba(&sf).unwrap();
        assert_eq!(model.samples_seen(), 0);
        assert_eq!(model.feature_count(), 0);
        model.learn(&sf, true).unwrap();
        let count = model.feature_count();
        let _ = model.predict_proba(&sf).unwrap();
        assert_eq!(model.feature_count(), count);
        assert_eq!(model.samples_seen(), 1);
    }

    #[test]
    fn classifier_non_finite_value_rejected() {
        let model = FtrlClassifier::new(FtrlConfig::default()).unwrap();
        assert!(SparseFeatures::from_sorted(vec![(0, f64::NAN)]).is_err());
        assert!(SparseFeatures::from_sorted(vec![(0, f64::INFINITY)]).is_err());
        let sf = SparseFeatures::from_sorted(vec![(0, 1.0)]).unwrap();
        assert!(model.predict_proba(&sf).is_ok());
    }

    #[test]
    fn classifier_empty_features_rejected() {
        let mut model = FtrlClassifier::new(FtrlConfig::default()).unwrap();
        let sf = SparseFeatures::new();
        assert!(model.predict_proba(&sf).is_err());
        assert!(model.learn(&sf, true).is_err());
    }

    #[test]
    fn classifier_reset_clears_state() {
        let mut model = FtrlClassifier::new(FtrlConfig::default()).unwrap();
        let sf = SparseFeatures::from_sorted(vec![(0, 1.0)]).unwrap();
        model.learn(&sf, true).unwrap();
        model.learn(&sf, false).unwrap();
        assert_eq!(model.samples_seen(), 2);
        assert!(model.feature_count() > 0);
        model.reset();
        assert_eq!(model.samples_seen(), 0);
        assert_eq!(model.feature_count(), 0);
        let p = model.predict_proba(&sf).unwrap();
        assert!((p - 0.5).abs() < 1e-12);
    }

    #[test]
    fn classifier_invalid_config_rejected() {
        assert!(
            FtrlClassifier::new(FtrlConfig {
                alpha: 0.0,
                ..FtrlConfig::default()
            })
            .is_err()
        );
        assert!(
            FtrlClassifier::new(FtrlConfig {
                beta: -0.1,
                ..FtrlConfig::default()
            })
            .is_err()
        );
        assert!(
            FtrlClassifier::new(FtrlConfig {
                l1: -1.0,
                ..FtrlConfig::default()
            })
            .is_err()
        );
        assert!(
            FtrlClassifier::new(FtrlConfig {
                l2: -1.0,
                ..FtrlConfig::default()
            })
            .is_err()
        );
        assert!(
            FtrlClassifier::new(FtrlConfig {
                alpha: f64::INFINITY,
                ..FtrlConfig::default()
            })
            .is_err()
        );
    }

    #[test]
    #[cfg(feature = "serde")]
    fn classifier_serde_roundtrip() {
        let mut model = FtrlClassifier::new(FtrlConfig {
            alpha: 0.3,
            beta: 0.5,
            l1: 0.1,
            l2: 0.2,
            max_features: Some(100),
            new_feature_policy: NewFeaturePolicy::Reject,
        })
        .unwrap();
        let sf = SparseFeatures::from_sorted(vec![(0, 1.0), (2, -1.0)]).unwrap();
        model.learn(&sf, true).unwrap();
        model.learn(&sf, false).unwrap();
        let json = serde_json::to_string(&model).unwrap();
        let restored: FtrlClassifier = serde_json::from_str(&json).unwrap();
        assert_eq!(restored.samples_seen(), model.samples_seen());
        assert_eq!(restored.feature_count(), model.feature_count());
        let p1 = model.predict_proba(&sf).unwrap();
        let p2 = restored.predict_proba(&sf).unwrap();
        assert!((p1 - p2).abs() < 1e-12);
    }

    #[test]
    fn predict_proba_in_range() {
        let mut model = FtrlClassifier::new(FtrlConfig {
            alpha: 0.5,
            beta: 1.0,
            l1: 0.0,
            l2: 0.0,
            max_features: None,
            new_feature_policy: NewFeaturePolicy::default(),
        })
        .unwrap();
        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(17);
        for _ in 0..200 {
            let x0 = rand::Rng::gen_range(&mut rng, -5.0..5.0);
            let x1 = rand::Rng::gen_range(&mut rng, -5.0..5.0);
            let y = x0 > 0.0;
            let sf = SparseFeatures::from_sorted(vec![(0, x0), (1, x1)]).unwrap();
            model.learn(&sf, y).unwrap();
            let p = model.predict_proba(&sf).unwrap();
            assert!(
                (0.0..=1.0).contains(&p),
                "probability must be in [0,1], got {p}"
            );
        }
    }

    #[test]
    fn learn_improves_accuracy() {
        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(21);
        // Generate a fixed test set.
        let test_set: Vec<(SparseFeatures, bool)> = (0..100)
            .map(|_| {
                let x0 = rand::Rng::gen_range(&mut rng, -2.0..2.0);
                let x1 = rand::Rng::gen_range(&mut rng, -1.0..1.0);
                let y = x0 + x1 > 0.0;
                (
                    SparseFeatures::from_sorted(vec![(0, x0), (1, x1)]).unwrap(),
                    y,
                )
            })
            .collect();

        let mut model = FtrlClassifier::new(FtrlConfig {
            alpha: 0.5,
            beta: 1.0,
            l1: 0.0,
            l2: 0.0,
            max_features: None,
            new_feature_policy: NewFeaturePolicy::default(),
        })
        .unwrap();

        // Accuracy before learning (always predicts 0.5 -> threshold 0.5 -> true).
        let acc_before: f64 = test_set
            .iter()
            .map(|(sf, y)| {
                let pred = model.predict(sf).unwrap();
                if pred == *y { 1.0 } else { 0.0 }
            })
            .sum::<f64>()
            / test_set.len() as f64;

        // Train on fresh data.
        for _ in 0..1000 {
            let x0 = rand::Rng::gen_range(&mut rng, -2.0..2.0);
            let x1 = rand::Rng::gen_range(&mut rng, -1.0..1.0);
            let y = x0 + x1 > 0.0;
            let sf = SparseFeatures::from_sorted(vec![(0, x0), (1, x1)]).unwrap();
            model.learn(&sf, y).unwrap();
        }

        let acc_after: f64 = test_set
            .iter()
            .map(|(sf, y)| {
                let pred = model.predict(sf).unwrap();
                if pred == *y { 1.0 } else { 0.0 }
            })
            .sum::<f64>()
            / test_set.len() as f64;

        assert!(
            acc_after > acc_before,
            "accuracy should improve: {acc_before} -> {acc_after}"
        );
    }

    #[test]
    fn classifier_weights_returns_nonzero_only() {
        let mut model = FtrlClassifier::new(FtrlConfig {
            alpha: 0.5,
            beta: 1.0,
            l1: 0.0,
            l2: 0.0,
            max_features: None,
            new_feature_policy: NewFeaturePolicy::default(),
        })
        .unwrap();
        let sf = SparseFeatures::from_sorted(vec![(0, 1.0), (1, 0.0001)]).unwrap();
        for _ in 0..50 {
            model.learn(&sf, true).unwrap();
        }
        let weights = model.weights();
        for &(_, w) in &weights {
            assert!(w != 0.0);
        }
    }

    #[test]
    fn classifier_multiple_features() {
        let mut model = FtrlClassifier::new(FtrlConfig {
            alpha: 0.5,
            beta: 1.0,
            l1: 0.0,
            l2: 0.0,
            max_features: None,
            new_feature_policy: NewFeaturePolicy::default(),
        })
        .unwrap();
        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(33);
        for _ in 0..1000 {
            let x0 = rand::Rng::gen_range(&mut rng, -2.0..2.0);
            let x1 = rand::Rng::gen_range(&mut rng, -2.0..2.0);
            let x2 = rand::Rng::gen_range(&mut rng, -1.0..1.0);
            // y = 1 if x0 + x1 > 0
            let y = x0 + x1 > 0.0;
            let sf = SparseFeatures::from_sorted(vec![(0, x0), (1, x1), (2, x2)]).unwrap();
            model.learn(&sf, y).unwrap();
        }
        let weights = model.weights();
        // Features 0 and 1 should have non-zero weights; feature 2 may or may not.
        assert!(weights.iter().any(|&(id, _)| id == 0));
        assert!(weights.iter().any(|&(id, _)| id == 1));
        // Verify prediction quality.
        let p_pos = model
            .predict_proba(
                &SparseFeatures::from_sorted(vec![(0, 3.0), (1, 3.0), (2, 0.0)]).unwrap(),
            )
            .unwrap();
        let p_neg = model
            .predict_proba(
                &SparseFeatures::from_sorted(vec![(0, -3.0), (1, -3.0), (2, 0.0)]).unwrap(),
            )
            .unwrap();
        assert!(p_pos > 0.8);
        assert!(p_neg < 0.2);
    }

    #[test]
    fn log_loss_converges() {
        // Average log loss should decrease over training.
        let mut model = FtrlClassifier::new(FtrlConfig {
            alpha: 0.5,
            beta: 1.0,
            l1: 0.0,
            l2: 0.0,
            max_features: None,
            new_feature_policy: NewFeaturePolicy::default(),
        })
        .unwrap();
        let loss_fn = crate::loss::log_loss::BinaryLogLoss::new();
        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(55);
        let mut first_loss = 0.0;
        let mut last_loss = 0.0;
        for i in 0..1000 {
            let x0 = rand::Rng::gen_range(&mut rng, -2.0..2.0);
            let x1 = rand::Rng::gen_range(&mut rng, -1.0..1.0);
            let y = x0 > 0.0;
            let sf = SparseFeatures::from_sorted(vec![(0, x0), (1, x1)]).unwrap();
            let p = model.predict_proba(&sf).unwrap();
            let loss = loss_fn.loss(p, y);
            if i < 20 {
                first_loss += loss;
            }
            if i >= 980 {
                last_loss += loss;
            }
            model.learn(&sf, y).unwrap();
        }
        assert!(last_loss < first_loss, "log loss should decrease");
    }

    // -----------------------------------------------------------------
    // FtrlClassifier: failure atomicity and overflow (ML-001/002)
    // -----------------------------------------------------------------

    #[test]
    fn classifier_overflow_does_not_mutate_state() {
        let mut model = FtrlClassifier::new(FtrlConfig {
            alpha: 0.1,
            beta: 1.0,
            l1: 0.0,
            l2: 0.0,
            max_features: None,
            new_feature_policy: NewFeaturePolicy::default(),
        })
        .unwrap();
        let sf = SparseFeatures::from_sorted(vec![(0, 1.0), (1, 1e300)]).unwrap();
        // Cold-start probability is 0.5, grad = 0.5 - 0.0 = 0.5.
        // g for feature 1 = 0.5 * 1e300 = 5e299 (finite), g^2 = 2.5e599 = inf.
        let result = model.learn(&sf, false);
        assert!(result.is_err(), "expected overflow error, got {result:?}");
        assert_eq!(model.samples_seen(), 0);
        assert_eq!(model.feature_count(), 0);
        assert!(model.params.is_empty());
        assert_eq!(model.intercept.z, 0.0);
        assert_eq!(model.intercept.n, 0.0);
    }

    #[test]
    fn classifier_partial_update_is_atomic() {
        let mut model = FtrlClassifier::new(FtrlConfig {
            alpha: 0.1,
            beta: 1.0,
            l1: 0.0,
            l2: 0.0,
            max_features: None,
            new_feature_policy: NewFeaturePolicy::default(),
        })
        .unwrap();
        let sf = SparseFeatures::from_sorted(vec![(0, 1.0), (1, 1e300)]).unwrap();
        assert!(model.learn(&sf, false).is_err());
        assert!(!model.params.contains_key(&0));
        assert!(!model.params.contains_key(&1));
        assert_eq!(model.samples_seen(), 0);
    }

    #[test]
    #[cfg(feature = "serde")]
    fn classifier_samples_seen_overflow_is_atomic() {
        let json = format!(
            "{{\"config\":{{\"alpha\":0.1,\"beta\":1.0,\"l1\":1.0,\"l2\":1.0,\"max_features\":null,\"new_feature_policy\":\"Reject\"}},\"params\":{{}},\"intercept\":{{\"z\":0.0,\"n\":0.0}},\"samples_seen\":{}}}",
            u64::MAX
        );
        let mut model: FtrlClassifier = serde_json::from_str(&json).unwrap();
        let sf = SparseFeatures::from_sorted(vec![(0, 1.0)]).unwrap();
        let result = model.learn(&sf, true);
        assert!(result.is_err(), "expected counter overflow");
        assert_eq!(model.samples_seen(), u64::MAX);
        assert_eq!(model.feature_count(), 0);
        assert_eq!(model.intercept.z, 0.0);
        assert_eq!(model.intercept.n, 0.0);
    }

    // -----------------------------------------------------------------
    // FtrlClassifier: max_features boundary (ML-003)
    // -----------------------------------------------------------------

    #[test]
    fn classifier_max_features_reject_at_limit() {
        let mut model = FtrlClassifier::new(FtrlConfig {
            alpha: 0.5,
            beta: 1.0,
            l1: 0.0,
            l2: 0.0,
            max_features: Some(2),
            new_feature_policy: NewFeaturePolicy::Reject,
        })
        .unwrap();
        let sf = SparseFeatures::from_sorted(vec![(0, 1.0), (1, 2.0)]).unwrap();
        model.learn(&sf, true).unwrap();
        assert_eq!(model.feature_count(), 2);
        let sf_new = SparseFeatures::from_sorted(vec![(2, 3.0)]).unwrap();
        assert!(model.learn(&sf_new, true).is_err());
        assert_eq!(model.feature_count(), 2);
        assert_eq!(model.samples_seen(), 1);
    }

    #[test]
    fn classifier_max_features_ignore_skips_new() {
        let mut model = FtrlClassifier::new(FtrlConfig {
            alpha: 0.5,
            beta: 1.0,
            l1: 0.0,
            l2: 0.0,
            max_features: Some(2),
            new_feature_policy: NewFeaturePolicy::Ignore,
        })
        .unwrap();
        let sf = SparseFeatures::from_sorted(vec![(0, 1.0), (1, 2.0)]).unwrap();
        model.learn(&sf, true).unwrap();
        let sf_mixed = SparseFeatures::from_sorted(vec![(0, 1.0), (2, 3.0)]).unwrap();
        model.learn(&sf_mixed, false).unwrap();
        assert_eq!(model.feature_count(), 2);
        assert!(!model.params.contains_key(&2));
        assert_eq!(model.samples_seen(), 2);
    }

    #[test]
    fn classifier_max_features_multi_new_prejudge() {
        let mut model = FtrlClassifier::new(FtrlConfig {
            alpha: 0.5,
            beta: 1.0,
            l1: 0.0,
            l2: 0.0,
            max_features: Some(2),
            new_feature_policy: NewFeaturePolicy::Reject,
        })
        .unwrap();
        let sf = SparseFeatures::from_sorted(vec![(0, 1.0), (1, 2.0), (2, 3.0)]).unwrap();
        assert!(model.learn(&sf, true).is_err());
        assert_eq!(model.feature_count(), 0);
        assert_eq!(model.samples_seen(), 0);
    }

    // -----------------------------------------------------------------
    // FtrlClassifier: serde validation (ML-004)
    // -----------------------------------------------------------------

    #[test]
    #[cfg(feature = "serde")]
    fn classifier_serde_rejects_negative_n() {
        let json = "{\"config\":{\"alpha\":0.1,\"beta\":1.0,\"l1\":1.0,\"l2\":1.0,\"max_features\":null,\"new_feature_policy\":\"Reject\"},\"params\":{\"0\":{\"z\":1.0,\"n\":-1.0}},\"intercept\":{\"z\":0.0,\"n\":0.0},\"samples_seen\":0}";
        let result: Result<FtrlClassifier, _> = serde_json::from_str(json);
        assert!(result.is_err(), "negative n must be rejected");
    }

    #[test]
    #[cfg(feature = "serde")]
    fn classifier_serde_rejects_invalid_config() {
        let json = "{\"config\":{\"alpha\":-1.0,\"beta\":1.0,\"l1\":1.0,\"l2\":1.0,\"max_features\":null,\"new_feature_policy\":\"Reject\"},\"params\":{},\"intercept\":{\"z\":0.0,\"n\":0.0},\"samples_seen\":0}";
        let result: Result<FtrlClassifier, _> = serde_json::from_str(json);
        assert!(result.is_err(), "invalid alpha must be rejected");
    }

    #[test]
    #[cfg(feature = "serde")]
    fn classifier_serde_accepts_missing_optional_fields() {
        let json = "{\"config\":{\"alpha\":0.1,\"beta\":1.0,\"l1\":1.0,\"l2\":1.0},\"params\":{},\"intercept\":{\"z\":0.0,\"n\":0.0},\"samples_seen\":0}";
        let model: FtrlClassifier = serde_json::from_str(json).unwrap();
        assert!(model.config().max_features.is_none());
        assert_eq!(model.config().new_feature_policy, NewFeaturePolicy::Reject);
    }
}