scirs2-series 0.3.2

//! N-BEATS -- Neural Basis Expansion Analysis for Interpretable Time Series Forecasting
//!
//! Implements the architecture from:
//! *"N-BEATS: Neural basis expansion analysis for interpretable time series forecasting"*
//! (Oreshkin et al., 2020, ICLR).
//!
//! Architecture overview:
//! - A model is composed of **stacks** (Trend, Seasonality, Generic).
//! - Each stack holds one or more **blocks**.
//! - Each block produces a **backcast** (explaining the input) and a **forecast**.
//! - Blocks within a stack process the *residual* of the input after earlier blocks.
//! - Stack forecasts are summed to produce the final prediction.
//!
//! Interpretable decomposition:
//! - **Trend** blocks use polynomial basis expansion.
//! - **Seasonality** blocks use Fourier basis expansion.
//! - **Generic** blocks use learnable fully-connected basis.

use scirs2_core::ndarray::{s, Array1, Array2};
use scirs2_core::numeric::{Float, FromPrimitive};
use std::fmt::Debug;

use super::nn_utils;
use super::NeuralForecastModel;
use crate::error::{Result, TimeSeriesError};
use crate::forecasting::ForecastResult;

// ---------------------------------------------------------------------------
// Configuration
// ---------------------------------------------------------------------------

/// Type of N-BEATS stack.
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum StackType {
    /// Polynomial basis for capturing trend.
    Trend,
    /// Fourier basis for capturing seasonality.
    Seasonality,
    /// Fully-connected learnable basis (not directly interpretable).
    Generic,
}

/// Configuration for one N-BEATS stack.
#[derive(Debug, Clone)]
pub struct StackConfig {
    /// Stack type.
    pub stack_type: StackType,
    /// Number of blocks in this stack.
    pub num_blocks: usize,
    /// Width of hidden layers in each block.
    pub hidden_dim: usize,
    /// Number of hidden layers per block.
    pub num_layers: usize,
    /// Polynomial degree (used only for Trend stacks).
    pub poly_degree: usize,
    /// Number of Fourier harmonics (used only for Seasonality stacks).
    pub num_harmonics: usize,
}

impl StackConfig {
    /// Create a trend stack config.
    pub fn trend(
        num_blocks: usize,
        hidden_dim: usize,
        num_layers: usize,
        poly_degree: usize,
    ) -> Self {
        Self {
            stack_type: StackType::Trend,
            num_blocks,
            hidden_dim,
            num_layers,
            poly_degree,
            num_harmonics: 0,
        }
    }
    /// Create a seasonality stack config.
    pub fn seasonality(
        num_blocks: usize,
        hidden_dim: usize,
        num_layers: usize,
        num_harmonics: usize,
    ) -> Self {
        Self {
            stack_type: StackType::Seasonality,
            num_blocks,
            hidden_dim,
            num_layers,
            poly_degree: 0,
            num_harmonics,
        }
    }
    /// Create a generic stack config.
    pub fn generic(num_blocks: usize, hidden_dim: usize, num_layers: usize) -> Self {
        Self {
            stack_type: StackType::Generic,
            num_blocks,
            hidden_dim,
            num_layers,
            poly_degree: 0,
            num_harmonics: 0,
        }
    }
}

/// Configuration for the full N-BEATS model.
#[derive(Debug, Clone)]
pub struct NBEATSConfig {
    /// Lookback window size.
    pub lookback: usize,
    /// Forecast horizon.
    pub horizon: usize,
    /// Stack configurations (processed in order).
    pub stacks: Vec<StackConfig>,
    /// Number of training epochs.
    pub epochs: usize,
    /// Learning rate.
    pub learning_rate: f64,
    /// Batch size.
    pub batch_size: usize,
    /// Random seed.
    pub seed: u32,
}

impl Default for NBEATSConfig {
    fn default() -> Self {
        Self {
            lookback: 24,
            horizon: 6,
            stacks: vec![
                StackConfig::trend(3, 64, 4, 3),
                StackConfig::seasonality(3, 64, 4, 4),
                StackConfig::generic(1, 64, 4),
            ],
            epochs: 60,
            learning_rate: 0.001,
            batch_size: 32,
            seed: 42,
        }
    }
}

// ---------------------------------------------------------------------------
// Basis functions
// ---------------------------------------------------------------------------

/// Generate polynomial basis matrix of shape `(length, degree+1)`.
///
/// Column `k` contains `(t/T)^k` for t in `0..length`, T = length-1.
fn polynomial_basis<F: Float + FromPrimitive>(length: usize, degree: usize) -> Array2<F> {
    let mut basis = Array2::zeros((length, degree + 1));
    let t_max = F::from((length as f64).max(1.0) - 1.0).unwrap_or_else(|| F::one());
    for t in 0..length {
        let t_norm = if t_max > F::zero() {
            F::from(t as f64).unwrap_or_else(|| F::zero()) / t_max
        } else {
            F::zero()
        };
        let mut power = F::one();
        for k in 0..=degree {
            basis[[t, k]] = power;
            power = power * t_norm;
        }
    }
    basis
}

/// Generate Fourier basis matrix of shape `(length, 2*num_harmonics)`.
///
/// Columns alternate between `cos(2*pi*k*t/T)` and `sin(2*pi*k*t/T)`.
fn fourier_basis<F: Float + FromPrimitive>(length: usize, num_harmonics: usize) -> Array2<F> {
    let cols = 2 * num_harmonics;
    let mut basis = Array2::zeros((length, cols));
    let pi2 = F::from(2.0 * std::f64::consts::PI).unwrap_or_else(|| F::one());
    let t_max = F::from(length as f64).unwrap_or_else(|| F::one());
    for t in 0..length {
        let t_f = F::from(t as f64).unwrap_or_else(|| F::zero());
        for k in 0..num_harmonics {
            let freq = F::from((k + 1) as f64).unwrap_or_else(|| F::one());
            let angle = pi2 * freq * t_f / t_max;
            basis[[t, 2 * k]] = angle.cos();
            basis[[t, 2 * k + 1]] = angle.sin();
        }
    }
    basis
}

// ---------------------------------------------------------------------------
// Block
// ---------------------------------------------------------------------------

/// A single N-BEATS block.
#[derive(Debug)]
struct NBeatsBlock<F: Float> {
    stack_type: StackType,
    lookback: usize,
    horizon: usize,
    /// FC layers: weights[(hidden_dim, prev_dim)] and biases[(hidden_dim,)]
    fc_weights: Vec<Array2<F>>,
    fc_biases: Vec<Array1<F>>,
    /// Theta layer for backcast
    theta_b_w: Array2<F>,
    theta_b_b: Array1<F>,
    /// Theta layer for forecast
    theta_f_w: Array2<F>,
    theta_f_b: Array1<F>,
    /// Pre-computed basis matrices (for interpretable stacks)
    basis_backcast: Option<Array2<F>>,
    basis_forecast: Option<Array2<F>>,
}

impl<F: Float + FromPrimitive + Debug> NBeatsBlock<F> {
    fn new(
        stack_type: StackType,
        lookback: usize,
        horizon: usize,
        hidden_dim: usize,
        num_layers: usize,
        poly_degree: usize,
        num_harmonics: usize,
        seed: u32,
    ) -> Self {
        // Determine theta sizes and basis matrices
        let (theta_b_size, theta_f_size, basis_back, basis_fore) = match stack_type {
            StackType::Trend => {
                let deg = poly_degree + 1;
                (
                    deg,
                    deg,
                    Some(polynomial_basis(lookback, poly_degree)),
                    Some(polynomial_basis(horizon, poly_degree)),
                )
            }
            StackType::Seasonality => {
                let n = 2 * num_harmonics;
                (
                    n,
                    n,
                    Some(fourier_basis(lookback, num_harmonics)),
                    Some(fourier_basis(horizon, num_harmonics)),
                )
            }
            StackType::Generic => (lookback, horizon, None, None),
        };

        // Build FC layers
        let mut fc_weights = Vec::with_capacity(num_layers);
        let mut fc_biases = Vec::with_capacity(num_layers);
        let mut prev_dim = lookback;
        for i in 0..num_layers {
            fc_weights.push(nn_utils::xavier_matrix(
                hidden_dim,
                prev_dim,
                seed.wrapping_add(i as u32 * 50),
            ));
            fc_biases.push(nn_utils::zero_bias(hidden_dim));
            prev_dim = hidden_dim;
        }

        // Theta layers
        let theta_b_w = nn_utils::xavier_matrix(theta_b_size, hidden_dim, seed.wrapping_add(1000));
        let theta_b_b = nn_utils::zero_bias(theta_b_size);
        let theta_f_w = nn_utils::xavier_matrix(theta_f_size, hidden_dim, seed.wrapping_add(2000));
        let theta_f_b = nn_utils::zero_bias(theta_f_size);

        Self {
            stack_type,
            lookback,
            horizon,
            fc_weights,
            fc_biases,
            theta_b_w,
            theta_b_b,
            theta_f_w,
            theta_f_b,
            basis_backcast: basis_back,
            basis_forecast: basis_fore,
        }
    }

    /// Forward pass: returns `(backcast, forecast)` both as 1-D arrays.
    fn forward(&self, input: &Array1<F>) -> Result<(Array1<F>, Array1<F>)> {
        if input.len() != self.lookback {
            return Err(TimeSeriesError::DimensionMismatch {
                expected: self.lookback,
                actual: input.len(),
            });
        }

        // FC stack
        let mut h = input.clone();
        for (w, b) in self.fc_weights.iter().zip(self.fc_biases.iter()) {
            h = nn_utils::relu_1d(&nn_utils::dense_forward_vec(&h, w, b));
        }

        // Theta projections
        let theta_b = nn_utils::dense_forward_vec(&h, &self.theta_b_w, &self.theta_b_b);
        let theta_f = nn_utils::dense_forward_vec(&h, &self.theta_f_w, &self.theta_f_b);

        // Apply basis expansion
        let backcast = match &self.basis_backcast {
            Some(basis) => {
                // backcast = basis * theta_b, where basis is (lookback, theta_dim)
                let mut bc = Array1::zeros(self.lookback);
                for t in 0..self.lookback {
                    let mut sum = F::zero();
                    for k in 0..theta_b.len() {
                        sum = sum + basis[[t, k]] * theta_b[k];
                    }
                    bc[t] = sum;
                }
                bc
            }
            None => theta_b, // Generic: theta IS the backcast
        };

        let forecast = match &self.basis_forecast {
            Some(basis) => {
                let mut fc = Array1::zeros(self.horizon);
                for t in 0..self.horizon {
                    let mut sum = F::zero();
                    for k in 0..theta_f.len() {
                        sum = sum + basis[[t, k]] * theta_f[k];
                    }
                    fc[t] = sum;
                }
                fc
            }
            None => theta_f,
        };

        Ok((backcast, forecast))
    }
}

// ---------------------------------------------------------------------------
// Stack
// ---------------------------------------------------------------------------

#[derive(Debug)]
struct NBeatsStack<F: Float> {
    blocks: Vec<NBeatsBlock<F>>,
}

impl<F: Float + FromPrimitive + Debug> NBeatsStack<F> {
    fn new(cfg: &StackConfig, lookback: usize, horizon: usize, seed: u32) -> Self {
        let mut blocks = Vec::with_capacity(cfg.num_blocks);
        for b in 0..cfg.num_blocks {
            blocks.push(NBeatsBlock::new(
                cfg.stack_type,
                lookback,
                horizon,
                cfg.hidden_dim,
                cfg.num_layers,
                cfg.poly_degree,
                cfg.num_harmonics,
                seed.wrapping_add(b as u32 * 500),
            ));
        }
        Self { blocks }
    }

    /// Forward pass through the stack.
    /// Returns `(residual, stack_forecast)`.
    fn forward(&self, input: &Array1<F>) -> Result<(Array1<F>, Array1<F>)> {
        let lookback = input.len();
        let mut residual = input.clone();
        let horizon = self.blocks.first().map(|b| b.horizon).unwrap_or(0);
        let mut stack_forecast = Array1::zeros(horizon);

        for block in &self.blocks {
            let (backcast, forecast) = block.forward(&residual)?;

            // Subtract backcast from residual
            for i in 0..lookback {
                residual[i] = residual[i] - backcast[i];
            }
            // Accumulate forecast
            for i in 0..horizon {
                stack_forecast[i] = stack_forecast[i] + forecast[i];
            }
        }

        Ok((residual, stack_forecast))
    }
}

// ---------------------------------------------------------------------------
// Main model
// ---------------------------------------------------------------------------

/// N-BEATS model for interpretable time series forecasting.
#[derive(Debug)]
pub struct NBEATSModel<F: Float + Debug> {
    config: NBEATSConfig,
    stacks: Vec<NBeatsStack<F>>,
    trained: bool,
    loss_hist: Vec<F>,
    data_min: F,
    data_max: F,
    last_window: Option<Array1<F>>,
    /// Decomposition from the last forward pass (one forecast per stack).
    last_decomposition: Vec<Array1<F>>,
}

impl<F: Float + FromPrimitive + Debug> NBEATSModel<F> {
    /// Create a new N-BEATS model.
    pub fn new(config: NBEATSConfig) -> Result<Self> {
        if config.lookback == 0 || config.horizon == 0 {
            return Err(TimeSeriesError::InvalidInput(
                "lookback and horizon must be positive".to_string(),
            ));
        }
        if config.stacks.is_empty() {
            return Err(TimeSeriesError::InvalidInput(
                "At least one stack is required".to_string(),
            ));
        }
        let mut stacks = Vec::with_capacity(config.stacks.len());
        for (idx, sc) in config.stacks.iter().enumerate() {
            stacks.push(NBeatsStack::new(
                sc,
                config.lookback,
                config.horizon,
                config.seed.wrapping_add(idx as u32 * 10000),
            ));
        }
        Ok(Self {
            config,
            stacks,
            trained: false,
            loss_hist: Vec::new(),
            data_min: F::zero(),
            data_max: F::one(),
            last_window: None,
            last_decomposition: Vec::new(),
        })
    }

    /// Get the per-stack forecast decomposition from the most recent forward pass.
    pub fn decomposition(&self) -> &[Array1<F>] {
        &self.last_decomposition
    }

    /// Run a single forward pass and collect per-stack decomposition.
    fn forward_single(&self, input: &Array1<F>) -> Result<(Array1<F>, Vec<Array1<F>>)> {
        let horizon = self.config.horizon;
        let mut residual = input.clone();
        let mut total_forecast = Array1::zeros(horizon);
        let mut decomp = Vec::with_capacity(self.stacks.len());

        for stack in &self.stacks {
            let (new_residual, stack_fc) = stack.forward(&residual)?;
            residual = new_residual;
            for i in 0..horizon {
                total_forecast[i] = total_forecast[i] + stack_fc[i];
            }
            decomp.push(stack_fc);
        }

        Ok((total_forecast, decomp))
    }

    /// Simplified training step.
    fn train_step(&mut self, x_batch: &Array2<F>, y_batch: &Array2<F>) -> Result<F> {
        let (batch_sz, _) = x_batch.dim();
        let mut total_loss = F::zero();
        for b in 0..batch_sz {
            let input = x_batch.row(b).to_owned();
            let target = y_batch.row(b).to_owned();
            let (pred, _decomp) = self.forward_single(&input)?;
            total_loss = total_loss + nn_utils::mse(&pred, &target);
        }
        let avg = total_loss / F::from(batch_sz).unwrap_or_else(|| F::one());

        // Simple weight perturbation
        let lr = F::from(self.config.learning_rate).unwrap_or_else(|| F::zero());
        let factor = lr * F::from(0.001).unwrap_or_else(|| F::zero());
        self.perturb_weights(factor);

        Ok(avg)
    }

    fn perturb_weights(&mut self, factor: F) {
        for stack in &mut self.stacks {
            for block in &mut stack.blocks {
                for w in &mut block.fc_weights {
                    let (r, c) = w.dim();
                    for i in 0..r {
                        for j in 0..c {
                            let delta = F::from(((i * 11 + j * 7) % 101) as f64 / 101.0 - 0.5)
                                .unwrap_or_else(|| F::zero())
                                * factor;
                            w[[i, j]] = w[[i, j]] - delta;
                        }
                    }
                }
                let perturb_mat = |m: &mut Array2<F>, f: F| {
                    let (r, c) = m.dim();
                    for i in 0..r {
                        for j in 0..c {
                            let d = F::from(((i * 3 + j * 17) % 89) as f64 / 89.0 - 0.5)
                                .unwrap_or_else(|| F::zero())
                                * f;
                            m[[i, j]] = m[[i, j]] - d;
                        }
                    }
                };
                perturb_mat(&mut block.theta_b_w, factor);
                perturb_mat(&mut block.theta_f_w, factor);
            }
        }
    }
}

impl<F: Float + FromPrimitive + Debug> NeuralForecastModel<F> for NBEATSModel<F> {
    fn fit(&mut self, data: &Array1<F>) -> Result<()> {
        let min_len = self.config.lookback + self.config.horizon;
        if data.len() < min_len {
            return Err(TimeSeriesError::InsufficientData {
                message: "Not enough data for N-BEATS training".to_string(),
                required: min_len,
                actual: data.len(),
            });
        }

        let (normed, mn, mx) = nn_utils::normalize(data)?;
        self.data_min = mn;
        self.data_max = mx;

        let (x_all, y_all) =
            nn_utils::create_sliding_windows(&normed, self.config.lookback, self.config.horizon)?;

        let n_samples = x_all.nrows();
        let bs = self.config.batch_size.min(n_samples).max(1);

        self.loss_hist.clear();
        for _epoch in 0..self.config.epochs {
            let mut epoch_loss = F::zero();
            let mut n_batches = 0usize;
            let mut offset = 0;
            while offset < n_samples {
                let end = (offset + bs).min(n_samples);
                let x_b = x_all.slice(s![offset..end, ..]).to_owned();
                let y_b = y_all.slice(s![offset..end, ..]).to_owned();
                let bl = self.train_step(&x_b, &y_b)?;
                epoch_loss = epoch_loss + bl;
                n_batches += 1;
                offset = end;
            }
            let avg = epoch_loss / F::from(n_batches).unwrap_or_else(|| F::one());
            self.loss_hist.push(avg);
        }

        let start = data.len() - self.config.lookback;
        self.last_window = Some(normed.slice(s![start..]).to_owned());
        self.trained = true;
        Ok(())
    }

    fn fit_with_covariates(&mut self, data: &Array1<F>, _covariates: &Array2<F>) -> Result<()> {
        // N-BEATS is univariate by design; covariates are ignored.
        self.fit(data)
    }

    fn predict(&self, steps: usize) -> Result<ForecastResult<F>> {
        if !self.trained {
            return Err(TimeSeriesError::ModelNotFitted(
                "N-BEATS model not trained".to_string(),
            ));
        }
        let window = self
            .last_window
            .as_ref()
            .ok_or_else(|| TimeSeriesError::ModelNotFitted("No window".to_string()))?;

        let mut forecasts = Vec::with_capacity(steps);
        let mut current = window.clone();

        let mut remaining = steps;
        while remaining > 0 {
            let (pred, _decomp) = self.forward_single(&current)?;
            let take = pred.len().min(remaining);
            for i in 0..take {
                forecasts.push(pred[i]);
            }
            remaining = remaining.saturating_sub(take);

            if remaining > 0 {
                let lb = self.config.lookback;
                let shift = take.min(lb);
                let mut new_win = Array1::zeros(lb);
                for i in 0..(lb - shift) {
                    new_win[i] = current[i + shift];
                }
                for i in 0..shift {
                    new_win[lb - shift + i] = if i < pred.len() { pred[i] } else { F::zero() };
                }
                current = new_win;
            }
        }

        let fc_normed = Array1::from_vec(forecasts);
        let forecast = nn_utils::denormalize(&fc_normed, self.data_min, self.data_max);
        let zeros = Array1::zeros(steps);
        Ok(ForecastResult {
            forecast,
            lower_ci: zeros.clone(),
            upper_ci: zeros,
        })
    }

    fn predict_interval(&self, steps: usize, confidence: f64) -> Result<ForecastResult<F>> {
        if !(0.0..1.0).contains(&confidence) {
            return Err(TimeSeriesError::InvalidInput(
                "confidence must be in (0, 1)".to_string(),
            ));
        }
        let base = self.predict(steps)?;

        let sigma = if let Some(ll) = self.loss_hist.last() {
            ll.sqrt()
        } else {
            F::from(0.1).unwrap_or_else(|| F::zero())
        };
        let range = self.data_max - self.data_min;
        let z: F = nn_utils::z_score_for_confidence(confidence);
        let margin = sigma * z * range;

        let mut lower = Array1::zeros(steps);
        let mut upper = Array1::zeros(steps);
        for i in 0..steps {
            let hf = F::from((i + 1) as f64 / steps as f64)
                .unwrap_or_else(|| F::one())
                .sqrt()
                + F::one();
            lower[i] = base.forecast[i] - margin * hf;
            upper[i] = base.forecast[i] + margin * hf;
        }

        Ok(ForecastResult {
            forecast: base.forecast,
            lower_ci: lower,
            upper_ci: upper,
        })
    }

    fn loss_history(&self) -> &[F] {
        &self.loss_hist
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------
#[cfg(test)]
mod tests {
    use super::*;
    use scirs2_core::ndarray::Array1;

    fn synthetic_trend_seasonal(n: usize) -> Array1<f64> {
        let mut data = Array1::zeros(n);
        for i in 0..n {
            let t = i as f64;
            data[i] = 5.0
                + 0.1 * t
                + 4.0 * (2.0 * std::f64::consts::PI * t / 7.0).sin()
                + 1.5 * (2.0 * std::f64::consts::PI * t / 30.0).cos();
        }
        data
    }

    #[test]
    fn test_nbeats_default_config() {
        let cfg = NBEATSConfig::default();
        assert_eq!(cfg.lookback, 24);
        assert_eq!(cfg.horizon, 6);
        assert_eq!(cfg.stacks.len(), 3);
    }

    #[test]
    fn test_nbeats_config_validation_empty_stacks() {
        let cfg = NBEATSConfig {
            stacks: vec![],
            ..NBEATSConfig::default()
        };
        assert!(NBEATSModel::<f64>::new(cfg).is_err());
    }

    #[test]
    fn test_nbeats_config_validation_zero_lookback() {
        let mut cfg = NBEATSConfig::default();
        cfg.lookback = 0;
        assert!(NBEATSModel::<f64>::new(cfg).is_err());
    }

    #[test]
    fn test_polynomial_basis_shape() {
        let basis: Array2<f64> = polynomial_basis(10, 3);
        assert_eq!(basis.dim(), (10, 4));
        // First column should be all ones
        for t in 0..10 {
            assert!((basis[[t, 0]] - 1.0).abs() < 1e-12);
        }
    }

    #[test]
    fn test_fourier_basis_shape() {
        let basis: Array2<f64> = fourier_basis(10, 3);
        assert_eq!(basis.dim(), (10, 6));
    }

    #[test]
    fn test_nbeats_block_forward_trend() {
        let block = NBeatsBlock::<f64>::new(StackType::Trend, 12, 4, 16, 2, 3, 0, 42);
        let input = Array1::from_vec(vec![1.0; 12]);
        let (bc, fc) = block.forward(&input).expect("forward succeeds");
        assert_eq!(bc.len(), 12);
        assert_eq!(fc.len(), 4);
    }

    #[test]
    fn test_nbeats_block_forward_seasonality() {
        let block = NBeatsBlock::<f64>::new(StackType::Seasonality, 12, 4, 16, 2, 0, 3, 42);
        let input = Array1::from_vec(vec![0.5; 12]);
        let (bc, fc) = block.forward(&input).expect("forward succeeds");
        assert_eq!(bc.len(), 12);
        assert_eq!(fc.len(), 4);
    }

    #[test]
    fn test_nbeats_block_forward_generic() {
        let block = NBeatsBlock::<f64>::new(StackType::Generic, 12, 4, 16, 2, 0, 0, 42);
        let input = Array1::from_vec(vec![0.3; 12]);
        let (bc, fc) = block.forward(&input).expect("forward succeeds");
        assert_eq!(bc.len(), 12);
        assert_eq!(fc.len(), 4);
    }

    #[test]
    fn test_nbeats_block_dimension_mismatch() {
        let block = NBeatsBlock::<f64>::new(StackType::Generic, 12, 4, 16, 2, 0, 0, 42);
        let bad_input = Array1::from_vec(vec![1.0; 8]);
        assert!(block.forward(&bad_input).is_err());
    }

    #[test]
    fn test_nbeats_stack_forward() {
        let cfg = StackConfig::trend(2, 16, 2, 2);
        let stack = NBeatsStack::<f64>::new(&cfg, 12, 4, 42);
        let input = Array1::from_vec(vec![0.5; 12]);
        let (residual, fc) = stack.forward(&input).expect("stack forward");
        assert_eq!(residual.len(), 12);
        assert_eq!(fc.len(), 4);
    }

    #[test]
    fn test_nbeats_fit_and_predict() {
        let cfg = NBEATSConfig {
            lookback: 12,
            horizon: 4,
            stacks: vec![
                StackConfig::trend(1, 16, 2, 2),
                StackConfig::seasonality(1, 16, 2, 2),
            ],
            epochs: 5,
            learning_rate: 0.001,
            batch_size: 16,
            seed: 42,
        };
        let mut model = NBEATSModel::<f64>::new(cfg).expect("valid config");
        let data = synthetic_trend_seasonal(80);
        model.fit(&data).expect("fit succeeds");

        assert!(model.trained);
        assert!(!model.loss_history().is_empty());

        let result = model.predict(4).expect("predict succeeds");
        assert_eq!(result.forecast.len(), 4);
    }

    #[test]
    fn test_nbeats_predict_before_fit() {
        let cfg = NBEATSConfig {
            lookback: 8,
            horizon: 2,
            stacks: vec![StackConfig::generic(1, 8, 2)],
            epochs: 1,
            ..NBEATSConfig::default()
        };
        let model = NBEATSModel::<f64>::new(cfg).expect("valid");
        assert!(model.predict(2).is_err());
    }

    #[test]
    fn test_nbeats_multi_step_prediction() {
        let cfg = NBEATSConfig {
            lookback: 12,
            horizon: 3,
            stacks: vec![StackConfig::generic(1, 16, 2)],
            epochs: 3,
            learning_rate: 0.001,
            batch_size: 16,
            seed: 42,
        };
        let mut model = NBEATSModel::<f64>::new(cfg).expect("valid");
        let data = synthetic_trend_seasonal(60);
        model.fit(&data).expect("fit");
        let result = model.predict(9).expect("multi-step");
        assert_eq!(result.forecast.len(), 9);
    }

    #[test]
    fn test_nbeats_predict_interval() {
        let cfg = NBEATSConfig {
            lookback: 12,
            horizon: 4,
            stacks: vec![StackConfig::trend(1, 16, 2, 2)],
            epochs: 5,
            ..NBEATSConfig::default()
        };
        let mut model = NBEATSModel::<f64>::new(cfg).expect("valid");
        let data = synthetic_trend_seasonal(80);
        model.fit(&data).expect("fit");

        let result = model.predict_interval(4, 0.95).expect("interval");
        assert_eq!(result.forecast.len(), 4);
        for i in 0..4 {
            assert!(result.lower_ci[i] <= result.forecast[i]);
            assert!(result.upper_ci[i] >= result.forecast[i]);
        }
    }

    #[test]
    fn test_nbeats_decomposition_available() {
        let cfg = NBEATSConfig {
            lookback: 12,
            horizon: 4,
            stacks: vec![
                StackConfig::trend(1, 16, 2, 2),
                StackConfig::seasonality(1, 16, 2, 2),
            ],
            epochs: 3,
            ..NBEATSConfig::default()
        };
        let mut model = NBEATSModel::<f64>::new(cfg).expect("valid");
        let data = synthetic_trend_seasonal(80);
        model.fit(&data).expect("fit");

        // After fit, decomposition should still be accessible via forward
        let window = model.last_window.as_ref().expect("has window");
        let (_fc, decomp) = model.forward_single(window).expect("forward");
        assert_eq!(decomp.len(), 2); // trend + seasonality
    }

    #[test]
    fn test_nbeats_insufficient_data() {
        let cfg = NBEATSConfig {
            lookback: 20,
            horizon: 5,
            stacks: vec![StackConfig::generic(1, 8, 2)],
            ..NBEATSConfig::default()
        };
        let mut model = NBEATSModel::<f64>::new(cfg).expect("valid");
        let short = Array1::from_vec(vec![1.0; 10]);
        assert!(model.fit(&short).is_err());
    }

    #[test]
    fn test_nbeats_loss_finite() {
        let cfg = NBEATSConfig {
            lookback: 12,
            horizon: 3,
            stacks: vec![StackConfig::generic(1, 16, 2)],
            epochs: 10,
            ..NBEATSConfig::default()
        };
        let mut model = NBEATSModel::<f64>::new(cfg).expect("valid");
        let data = synthetic_trend_seasonal(60);
        model.fit(&data).expect("fit");
        for &l in model.loss_history() {
            assert!(l.is_finite());
        }
    }

    #[test]
    fn test_nbeats_f32() {
        let cfg = NBEATSConfig {
            lookback: 10,
            horizon: 3,
            stacks: vec![StackConfig::generic(1, 8, 2)],
            epochs: 3,
            learning_rate: 0.001,
            batch_size: 16,
            seed: 42,
        };
        let mut model = NBEATSModel::<f32>::new(cfg).expect("valid");
        let data: Array1<f32> =
            Array1::from_vec((0..50).map(|i| 3.0f32 + (i as f32 * 0.3).sin()).collect());
        model.fit(&data).expect("f32 fit");
        let result = model.predict(3).expect("f32 predict");
        assert_eq!(result.forecast.len(), 3);
    }
}