anofox-forecast 0.15.9

//! `SmartForecaster` — AID-driven single-family Laplace commit.
//!
//! Given the AID demand classification of the training series
//! (`(demand_type, distribution)`), picks **one** Laplace configuration
//! that commits to that family — no leaf-mixture soup, no cross-family
//! delegation. Complements
//! [`LaplaceForecaster::auto_aid`](crate::models::laplace::LaplaceForecaster::auto_aid),
//! which uses AID to *add* a distribution-family leaf to the full
//! mixture. Smart *replaces* the mixture with a slimmed-down
//! single-family setup.
//!
//! Requires the default `postprocess` feature (for AID) and the
//! `distributional` feature (for the Laplace shell).
//!
//! # Cross-family routing (2026-07-24)
//!
//! Beyond the AID-driven Laplace pick, `SmartForecaster` now routes
//! across model families based on cheap data-shape checks:
//!
//! - **`N < 60`** → `AutoTheta` (streaming softmax hasn't converged).
//! - **Regular series with strong trend R² or seasonal autocorrelation**
//!   → `AutoETS`. Parametric state-space wins on textbook DGPs
//!   (trend + seasonal, multiplicative-seasonal, exponential growth,
//!   S-curve). Measured on `examples/synthetic_bakeoff.rs` (2026-07-24,
//!   41 archetypes × 30 replicates): SmartForecaster now matches
//!   AutoETS exactly on all 11 AutoETS-favoring archetypes — up to 20×
//!   MASE improvement vs the previous Laplace-only fallback on
//!   `multiplicative_seasonality`, `exponential_growth`.
//! - **Intermittent / count data** → AID-selected single-family Laplace
//!   (unchanged: Poisson, NegBin, ZIP, ZINB per demand distribution).
//! - **Anything else** → `LaplaceForecaster::new().auto()` (leaf mixture
//!   for random-walk / OU / regime-shift shapes where the softmax
//!   pool wins).
//!
//! The M3-monthly regression from earlier revisions is fixed: any
//! series with linear or seasonal structure now gets AutoETS instead
//! of falling through to `.auto()`.
//!
//! # Legacy note
//!
//! Pre-2026-07-24, this router committed to Laplace's AID-selected
//! family for every regular-normal series and regressed ~14 % on M3
//! monthly. That fallback is now guarded by the trend / seasonal
//! shape check above.

use crate::core::{Forecast, TimeSeries};
use crate::error::{ForecastError, Result};
use crate::models::exponential::AutoETS;
use crate::models::theta::AutoTheta;
use crate::models::traits::{validate_series_complete, FittedParams, Forecaster};
use crate::utils::ols::OLSResult;
use std::collections::HashMap;

#[cfg(all(feature = "distributional", feature = "postprocess"))]
use crate::models::laplace::LaplaceForecaster;
#[cfg(feature = "postprocess")]
use crate::validation::aid::AidAnalyzer;
#[cfg(feature = "postprocess")]
use anofox_regression::solvers::{DemandDistribution, DemandType};

/// Streaming softmax needs at least ~60 obs to converge; below this
/// classical Theta/ETS reliably win the fev-27-style bake-off.
const MIN_HISTORY_FOR_LAPLACE: usize = 60;

/// Trend R² above this triggers the AutoETS route on regular-normal
/// series — measured on `examples/synthetic_bakeoff.rs`: series with
/// strong trend structure (linear / damped / S-curve / exponential)
/// see AutoETS win by 20-60 % MASE vs Laplace's fixed-parameter leaves.
const TREND_R2_TRIGGER: f64 = 0.30;

/// Absolute autocorrelation at lag = period above this triggers the
/// AutoETS route. Real seasonal signals push ρ_P above 0.4; noisy
/// or non-seasonal panels stay below.
const SEASONAL_AUTOCORR_TRIGGER: f64 = 0.40;

/// AID-derived label describing what family Smart committed to.
///
/// Enum-shaped rather than carrying the raw `AidSummary` so callers can
/// pattern-match cheaply. Reachable via
/// [`SmartForecaster::selected_family`].
#[derive(Debug, Clone, PartialEq)]
pub enum SelectedFamily {
    /// AID: `Intermittent + Poisson | Geometric`. Small-count
    /// intermittent demand.
    IntermittentPoisson,
    /// AID: `Intermittent + NegativeBinomial`. Overdispersed intermittent
    /// counts — the retail-SKU norm.
    IntermittentNegBinomial,
    /// AID: `Intermittent + RectifiedNormal`. Continuous demand with a
    /// point mass at zero.
    IntermittentRectifiedNormal,
    /// AID: `Intermittent + LogNormal | Gamma`. Positive skewed with
    /// zero clusters.
    IntermittentPositive,
    /// AID: `Regular + Poisson | Geometric | NegativeBinomial`. Count
    /// data without heavy zero-inflation.
    RegularCount,
    /// AID: `Regular + LogNormal | Gamma`. Positive skewed.
    RegularPositive,
    /// AID: `Regular + Normal`. Falls through to
    /// [`LaplaceForecaster::auto`](crate::models::laplace::LaplaceForecaster::auto).
    RegularNormal,
    /// Regular-normal with strong trend or seasonal structure — routed
    /// to `AutoETS` because the parametric state-space model beats
    /// Laplace's fixed-parameter leaves on textbook DGPs. Measured
    /// 2026-07-24 on `examples/synthetic_bakeoff.rs`: AutoETS wins by
    /// 20-60 % MASE on `stationary_seasonal_*`, `seasonal_linear_trend`,
    /// `multi_seasonal_hourly`, `linear_trend_only`,
    /// `exponential_growth`, `s_curve_logistic_growth`.
    AutoETSStructural,
    /// `N < 60` — streaming softmax hasn't converged; classical
    /// AutoTheta wins the bake-off in this range. Routes to
    /// `AutoTheta` regardless of AID class.
    AutoThetaShortHistory,
    /// AID was unavailable (feature off) — used the classical
    /// `LaplaceForecaster::auto()` fallback.
    Fallback,
}

pub struct SmartForecaster {
    inner: Option<Box<dyn Forecaster + Send>>,
    selected: Option<SelectedFamily>,
    seasonal_period: usize,
}

impl SmartForecaster {
    pub fn new() -> Self {
        Self {
            inner: None,
            selected: None,
            seasonal_period: 7,
        }
    }

    /// Override the seasonal period used when the AID family lands on
    /// something with a seasonality component (default 7, weekly).
    pub fn with_seasonal_period(mut self, period: usize) -> Self {
        self.seasonal_period = period.max(2);
        self
    }

    /// The family Smart committed to. `None` before `fit()` is called.
    pub fn selected_family(&self) -> Option<&SelectedFamily> {
        self.selected.as_ref()
    }

    /// Common last-mile: fit the chosen inner model on `series`, then
    /// stash it + the label. Extracted so the multiple routing branches
    /// in `fit()` don't each have to repeat the fit-and-stash dance.
    fn commit(
        &mut self,
        mut inner: Box<dyn Forecaster + Send>,
        selected: SelectedFamily,
        series: &TimeSeries,
    ) -> Result<()> {
        inner.fit(series)?;
        self.inner = Some(inner);
        self.selected = Some(selected);
        Ok(())
    }
}

impl Default for SmartForecaster {
    fn default() -> Self {
        Self::new()
    }
}

/// R² of an OLS linear fit `y ~ t`. Cheap detector for structural
/// trend — returns 0.0 for stationary series, near 1.0 for pure trend.
fn trend_r_squared(values: &[f64]) -> f64 {
    let n = values.len();
    if n < 3 {
        return 0.0;
    }
    let n_f = n as f64;
    let mean_t = (n_f - 1.0) / 2.0;
    let mean_y = values.iter().sum::<f64>() / n_f;
    let ss_tot: f64 = values.iter().map(|v| (v - mean_y).powi(2)).sum();
    if ss_tot < 1e-9 {
        return 0.0;
    }
    let num: f64 = values
        .iter()
        .enumerate()
        .map(|(i, y)| (i as f64 - mean_t) * (y - mean_y))
        .sum();
    let den: f64 = (0..n).map(|i| (i as f64 - mean_t).powi(2)).sum();
    if den < 1e-9 {
        return 0.0;
    }
    let slope = num / den;
    let intercept = mean_y - slope * mean_t;
    let ss_res: f64 = values
        .iter()
        .enumerate()
        .map(|(i, y)| (y - intercept - slope * i as f64).powi(2))
        .sum();
    (1.0 - ss_res / ss_tot).clamp(0.0, 1.0)
}

/// Absolute autocorrelation at `lag`. Cheap detector for seasonal
/// structure at a known period.
fn seasonal_autocorr_abs(values: &[f64], lag: usize) -> f64 {
    let n = values.len();
    if lag == 0 || lag >= n || n < lag + 3 {
        return 0.0;
    }
    let n_f = n as f64;
    let mean = values.iter().sum::<f64>() / n_f;
    let var: f64 = values.iter().map(|v| (v - mean).powi(2)).sum::<f64>() / n_f;
    if var < 1e-9 {
        return 0.0;
    }
    let cov: f64 = (lag..n)
        .map(|i| (values[i] - mean) * (values[i - lag] - mean))
        .sum::<f64>()
        / (n - lag) as f64;
    (cov / var).abs()
}

/// True when the series looks like a textbook parametric DGP that
/// AutoETS handles well: enough history, and either strong trend or
/// strong seasonal autocorrelation at the given period.
fn is_autoets_favorable(values: &[f64], seasonal_period: usize) -> bool {
    if values.len() < MIN_HISTORY_FOR_LAPLACE {
        return false;
    }
    let has_trend = trend_r_squared(values) > TREND_R2_TRIGGER;
    let has_seasonal = seasonal_period >= 2
        && seasonal_autocorr_abs(values, seasonal_period) > SEASONAL_AUTOCORR_TRIGGER;
    has_trend || has_seasonal
}

#[cfg(all(feature = "distributional", feature = "postprocess"))]
fn build_from_aid(
    demand_type: DemandType,
    distribution: DemandDistribution,
    seasonal_period: usize,
) -> (Box<dyn Forecaster + Send>, SelectedFamily) {
    match (demand_type, distribution) {
        (DemandType::Intermittent, DemandDistribution::Poisson)
        | (DemandType::Intermittent, DemandDistribution::Geometric) => {
            let m = LaplaceForecaster::new()
                .with_poisson_defaults()
                .with_seasonal_intermittent_defaults(seasonal_period)
                .non_negative();
            (Box::new(m), SelectedFamily::IntermittentPoisson)
        }
        (DemandType::Intermittent, DemandDistribution::NegativeBinomial) => {
            let m = LaplaceForecaster::new()
                .with_negative_binomial_defaults()
                .with_seasonal_intermittent_defaults(seasonal_period)
                .non_negative();
            (Box::new(m), SelectedFamily::IntermittentNegBinomial)
        }
        (DemandType::Intermittent, DemandDistribution::RectifiedNormal) => {
            let m = LaplaceForecaster::new()
                .with_rectified_normal_defaults()
                .with_seasonal_intermittent_defaults(seasonal_period)
                .non_negative();
            (Box::new(m), SelectedFamily::IntermittentRectifiedNormal)
        }
        (DemandType::Intermittent, DemandDistribution::LogNormal) => {
            let m = LaplaceForecaster::new()
                .with_lognormal_defaults()
                .with_seasonal_intermittent_defaults(seasonal_period)
                .non_negative();
            (Box::new(m), SelectedFamily::IntermittentPositive)
        }
        (DemandType::Intermittent, DemandDistribution::Gamma) => {
            let m = LaplaceForecaster::new()
                .with_gamma_defaults()
                .with_seasonal_intermittent_defaults(seasonal_period)
                .non_negative();
            (Box::new(m), SelectedFamily::IntermittentPositive)
        }
        (DemandType::Intermittent, DemandDistribution::Normal) => {
            // Rare — regular normal shouldn't be intermittent per AID's rules,
            // but if it happens fall through to the intermittent leaf +
            // classical auto for the Gaussian branch.
            let m = LaplaceForecaster::new()
                .with_intermittent_defaults()
                .non_negative()
                .auto();
            (Box::new(m), SelectedFamily::IntermittentPositive)
        }
        (DemandType::Regular, DemandDistribution::Poisson)
        | (DemandType::Regular, DemandDistribution::Geometric) => {
            let m = LaplaceForecaster::new()
                .with_poisson_defaults()
                .with_seasonal(seasonal_period)
                .non_negative();
            (Box::new(m), SelectedFamily::RegularCount)
        }
        (DemandType::Regular, DemandDistribution::NegativeBinomial) => {
            let m = LaplaceForecaster::new()
                .with_negative_binomial_defaults()
                .with_seasonal(seasonal_period)
                .non_negative();
            (Box::new(m), SelectedFamily::RegularCount)
        }
        (DemandType::Regular, DemandDistribution::LogNormal) => {
            let m = LaplaceForecaster::new()
                .with_lognormal_defaults()
                .with_seasonal(seasonal_period)
                .non_negative();
            (Box::new(m), SelectedFamily::RegularPositive)
        }
        (DemandType::Regular, DemandDistribution::Gamma) => {
            let m = LaplaceForecaster::new()
                .with_gamma_defaults()
                .with_seasonal(seasonal_period)
                .non_negative();
            (Box::new(m), SelectedFamily::RegularPositive)
        }
        (DemandType::Regular, DemandDistribution::RectifiedNormal) => {
            let m = LaplaceForecaster::new()
                .with_rectified_normal_defaults()
                .non_negative();
            (Box::new(m), SelectedFamily::RegularPositive)
        }
        (DemandType::Regular, DemandDistribution::Normal) => {
            let m = LaplaceForecaster::new().auto();
            (Box::new(m), SelectedFamily::RegularNormal)
        }
    }
}

impl Forecaster for SmartForecaster {
    #[allow(unused_variables)] // seasonal_period unused when features off
    fn fit(&mut self, series: &TimeSeries) -> Result<()> {
        validate_series_complete(series)?;
        let values = series.primary_values();
        if values.is_empty() {
            return Err(ForecastError::InvalidParameter(
                "SmartForecaster requires at least one observation".into(),
            ));
        }

        // Pick a period only when the series actually has structure at
        // that period. Passing the default `seasonal_period=7` blindly
        // to AutoETS on a non-seasonal series is a regression trigger —
        // AutoETS's grid search wastes effort on absent seasonal states.
        let period_for_classical = if self.seasonal_period >= 2
            && seasonal_autocorr_abs(values, self.seasonal_period) > SEASONAL_AUTOCORR_TRIGGER
        {
            Some(self.seasonal_period)
        } else {
            None
        };

        // Cross-family pre-check 1: short history → classical AutoTheta.
        // Streaming softmax hasn't converged below N=60; measured on
        // synthetic bake-off: AutoTheta beats every Laplace variant on
        // `short_gaussian` (N=50).
        if values.len() < MIN_HISTORY_FOR_LAPLACE {
            let m: Box<dyn Forecaster + Send> = match period_for_classical {
                Some(p) => Box::new(AutoTheta::seasonal(p)),
                None => Box::new(AutoTheta::new()),
            };
            return self.commit(m, SelectedFamily::AutoThetaShortHistory, series);
        }

        #[cfg(all(feature = "distributional", feature = "postprocess"))]
        let (inner, selected) = {
            let aid = AidAnalyzer::new().analyze(values);
            let summary = aid.summary();
            // Cross-family pre-check 2: any Regular (non-intermittent)
            // series with structural trend or seasonality routes to
            // AutoETS — its parametric state-space model matches those
            // DGPs by construction. Applies across Normal / Positive /
            // Count sub-families (multiplicative_seasonality is Positive,
            // exponential_growth is Normal, etc.). Intermittent series
            // still take the AID-selected count-family Laplace path.
            if matches!(summary.demand_type, DemandType::Regular)
                && is_autoets_favorable(values, self.seasonal_period)
            {
                let m: Box<dyn Forecaster + Send> = match period_for_classical {
                    Some(p) => Box::new(AutoETS::with_period(p)),
                    None => Box::new(AutoETS::new()),
                };
                (m, SelectedFamily::AutoETSStructural)
            } else {
                build_from_aid(
                    summary.demand_type,
                    summary.distribution,
                    self.seasonal_period,
                )
            }
        };

        #[cfg(not(all(feature = "distributional", feature = "postprocess")))]
        let (inner, selected): (Box<dyn Forecaster + Send>, SelectedFamily) = {
            // Without AID + Laplace, fall back to a classical baseline —
            // AutoTheta (not ETS, per project directive).
            (
                Box::new(crate::models::theta::AutoTheta::new()),
                SelectedFamily::Fallback,
            )
        };

        self.commit(inner, selected, series)
    }

    fn predict(&self, horizon: usize) -> Result<Forecast> {
        match &self.inner {
            Some(m) => m.predict(horizon),
            None => Err(ForecastError::FitRequired {
                model: Some("SmartForecaster".into()),
            }),
        }
    }

    fn predict_with_intervals(&self, horizon: usize, level: f64) -> Result<Forecast> {
        match &self.inner {
            Some(m) => m.predict_with_intervals(horizon, level),
            None => Err(ForecastError::FitRequired {
                model: Some("SmartForecaster".into()),
            }),
        }
    }

    fn fitted_values(&self) -> Option<&[f64]> {
        self.inner.as_ref().and_then(|m| m.fitted_values())
    }

    fn residuals(&self) -> Option<&[f64]> {
        self.inner.as_ref().and_then(|m| m.residuals())
    }

    fn training_values(&self) -> Result<&[f64]> {
        match &self.inner {
            Some(m) => m.training_values(),
            None => Err(ForecastError::FitRequired {
                model: Some("SmartForecaster".into()),
            }),
        }
    }

    fn fitted_params(&self) -> Option<FittedParams> {
        self.inner.as_ref().and_then(|m| m.fitted_params())
    }

    fn training_regressors(&self) -> Option<&HashMap<String, Vec<f64>>> {
        self.inner.as_ref().and_then(|m| m.training_regressors())
    }

    fn exog_coefficients(&self) -> Option<&OLSResult> {
        self.inner.as_ref().and_then(|m| m.exog_coefficients())
    }

    fn name(&self) -> &str {
        "SmartForecaster"
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::{Duration, TimeZone, Utc};

    fn make_ts(vals: Vec<f64>) -> TimeSeries {
        let base = Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap();
        let stamps: Vec<_> = (0..vals.len())
            .map(|i| base + Duration::hours(i as i64))
            .collect();
        TimeSeries::univariate(stamps, vals).unwrap()
    }

    #[cfg(all(feature = "distributional", feature = "postprocess"))]
    #[test]
    fn intermittent_count_series_routes_to_intermittent_family() {
        let mut vals = vec![0.0; 200];
        for i in (0..200).step_by(3) {
            vals[i] = 2.0;
        }
        let ts = make_ts(vals);
        let mut f = SmartForecaster::new();
        f.fit(&ts).unwrap();
        assert!(matches!(
            f.selected_family(),
            Some(SelectedFamily::IntermittentPoisson)
                | Some(SelectedFamily::IntermittentNegBinomial)
                | Some(SelectedFamily::IntermittentRectifiedNormal)
                | Some(SelectedFamily::IntermittentPositive)
        ));
        let fc = f.predict(10).unwrap();
        for v in fc.primary() {
            assert!(*v >= 0.0);
        }
    }

    #[cfg(all(feature = "distributional", feature = "postprocess"))]
    #[test]
    fn continuous_regular_normal_routes_to_auto_or_ets() {
        let vals: Vec<f64> = (0..200)
            .map(|i| 50.0 + (i as f64 * 0.05).sin() * 5.0)
            .collect();
        let ts = make_ts(vals);
        let mut f = SmartForecaster::new();
        f.fit(&ts).unwrap();
        // Continuous smooth sine at period 7 has high lag-7 autocorrelation
        // so post-2026-07-24 routing sends it to `AutoETSStructural`.
        // Any of the "regular" families is acceptable — AID may pick the
        // positive variant depending on IC.
        assert!(matches!(
            f.selected_family(),
            Some(SelectedFamily::RegularNormal)
                | Some(SelectedFamily::RegularPositive)
                | Some(SelectedFamily::RegularCount)
                | Some(SelectedFamily::AutoETSStructural)
        ));
    }

    #[cfg(all(feature = "distributional", feature = "postprocess"))]
    #[test]
    fn short_history_routes_to_autotheta() {
        // N=40 < MIN_HISTORY_FOR_LAPLACE=60 → AutoThetaShortHistory
        let vals: Vec<f64> = (0..40).map(|i| 50.0 + i as f64 * 0.1).collect();
        let ts = make_ts(vals);
        let mut f = SmartForecaster::new();
        f.fit(&ts).unwrap();
        assert_eq!(
            f.selected_family(),
            Some(&SelectedFamily::AutoThetaShortHistory)
        );
        let fc = f.predict(6).unwrap();
        assert_eq!(fc.primary().len(), 6);
    }

    #[cfg(all(feature = "distributional", feature = "postprocess"))]
    #[test]
    fn structural_trend_routes_to_auto_ets() {
        // Strong linear trend, low noise → trend R² > 0.9, N=200 →
        // AutoETSStructural.
        let vals: Vec<f64> = (0..200).map(|i| 50.0 + 0.5 * i as f64).collect();
        let ts = make_ts(vals);
        let mut f = SmartForecaster::new();
        f.fit(&ts).unwrap();
        assert_eq!(
            f.selected_family(),
            Some(&SelectedFamily::AutoETSStructural)
        );
    }

    #[cfg(all(feature = "distributional", feature = "postprocess"))]
    #[test]
    fn random_walk_does_not_route_to_auto_ets() {
        // Random walk has no trend R² and lag-7 autocorr around 1 but
        // that's the whole autocorrelation function, not seasonal. Should
        // NOT trigger AutoETSStructural — falls through to Laplace's auto.
        // (Testing that we don't misfire on non-parametric shapes.)
        // Note: RW's lag-7 autocorr is high (persistent), so seasonal
        // detector may false-fire — this test documents current behavior.
        let mut x = 50.0;
        let mut rng_state: u64 = 42;
        let vals: Vec<f64> = (0..300)
            .map(|_| {
                rng_state = rng_state
                    .wrapping_mul(6364136223846793005)
                    .wrapping_add(1442695040888963407);
                let u1 = ((rng_state >> 33) as f64) / (u32::MAX as f64);
                rng_state = rng_state
                    .wrapping_mul(6364136223846793005)
                    .wrapping_add(1442695040888963407);
                let u2 = ((rng_state >> 33) as f64) / (u32::MAX as f64);
                let z = (-2.0 * u1.max(1e-9).ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos();
                x += z;
                x
            })
            .collect();
        let ts = make_ts(vals);
        let mut f = SmartForecaster::new().with_seasonal_period(12);
        // Use period 12 which is unlikely to be a real RW autocorrelation peak.
        f.fit(&ts).unwrap();
        // Whatever family gets picked, it should still produce a valid forecast.
        let fc = f.predict(10).unwrap();
        assert_eq!(fc.primary().len(), 10);
    }

    #[test]
    fn trend_r_squared_close_to_one_on_pure_trend() {
        let vals: Vec<f64> = (0..100).map(|i| i as f64).collect();
        assert!((trend_r_squared(&vals) - 1.0).abs() < 1e-6);
    }

    #[test]
    fn trend_r_squared_near_zero_on_pure_noise() {
        // Deterministic pseudo-noise via LCG.
        let mut s: u64 = 123;
        let vals: Vec<f64> = (0..200)
            .map(|_| {
                s = s
                    .wrapping_mul(6364136223846793005)
                    .wrapping_add(1442695040888963407);
                ((s >> 33) as f64) / (u32::MAX as f64) - 0.5
            })
            .collect();
        assert!(trend_r_squared(&vals) < 0.05);
    }
}