scirs2-series 0.3.3

//! Structural time series models (Harvey's Basic Structural Model framework).
//!
//! Implements:
//! - **Local Level Model** (random walk + noise): Harvey (1989) §3.2
//! - **Local Linear Trend Model**: Harvey (1989) §3.3
//! - **Basic Structural Model** (BSM) with trigonometric seasonality: Harvey §3.4
//! - **Structural decomposition** into trend, seasonal, and irregular components
//! - **Forecasting** from fitted BSM
//!
//! All parameter estimation uses MLE via prediction-error decomposition (Kalman
//! filter log-likelihood), optimised with a simple bounded Nelder-Mead search.
//!
//! ## Model overview
//!
//! ```text
//! y_t = μ_t + γ_t + ε_t,    ε_t  ~ N(0, σ²_ε)
//! μ_{t+1} = μ_t + β_t + η_t, η_t  ~ N(0, σ²_η)   (level)
//! β_{t+1} = β_t  + ζ_t,      ζ_t  ~ N(0, σ²_ζ)   (slope)
//! γ_{t+1} = Σ_j [cos(λ_j)γ_{j,t} + sin(λ_j)γ*_{j,t}] + ω_{j,t}
//! ```

use crate::error::{Result, TimeSeriesError};
use std::f64::consts::PI;

// ---------------------------------------------------------------------------
// Internal flat-matrix algebra (row-major, no external BLAS)
// ---------------------------------------------------------------------------

/// A simple flat row-major matrix representation.
#[derive(Clone, Debug)]
struct Mat {
    rows: usize,
    cols: usize,
    data: Vec<f64>,
}

impl Mat {
    fn zeros(rows: usize, cols: usize) -> Self {
        Mat { rows, cols, data: vec![0.0; rows * cols] }
    }

    fn eye(n: usize) -> Self {
        let mut m = Self::zeros(n, n);
        for i in 0..n { m.set(i, i, 1.0); }
        m
    }

    #[inline]
    fn get(&self, r: usize, c: usize) -> f64 { self.data[r * self.cols + c] }

    #[inline]
    fn set(&mut self, r: usize, c: usize, v: f64) { self.data[r * self.cols + c] = v; }

    #[inline]
    fn add_assign(&mut self, r: usize, c: usize, v: f64) { self.data[r * self.cols + c] += v; }

    /// Matrix multiplication: self * rhs.
    fn mul(&self, rhs: &Mat) -> Self {
        assert_eq!(self.cols, rhs.rows, "mat_mul dimension mismatch");
        let mut out = Mat::zeros(self.rows, rhs.cols);
        for i in 0..self.rows {
            for k in 0..self.cols {
                let a = self.get(i, k);
                if a == 0.0 { continue; }
                for j in 0..rhs.cols {
                    let prev = out.get(i, j);
                    out.set(i, j, prev + a * rhs.get(k, j));
                }
            }
        }
        out
    }

    /// Transpose.
    fn t(&self) -> Self {
        let mut out = Mat::zeros(self.cols, self.rows);
        for i in 0..self.rows {
            for j in 0..self.cols {
                out.set(j, i, self.get(i, j));
            }
        }
        out
    }

    /// Element-wise addition.
    fn add(&self, rhs: &Mat) -> Self {
        assert_eq!(self.rows, rhs.rows);
        assert_eq!(self.cols, rhs.cols);
        let data: Vec<f64> = self.data.iter().zip(&rhs.data).map(|(a, b)| a + b).collect();
        Mat { rows: self.rows, cols: self.cols, data }
    }

    /// Element-wise subtraction.
    fn sub(&self, rhs: &Mat) -> Self {
        assert_eq!(self.rows, rhs.rows);
        assert_eq!(self.cols, rhs.cols);
        let data: Vec<f64> = self.data.iter().zip(&rhs.data).map(|(a, b)| a - b).collect();
        Mat { rows: self.rows, cols: self.cols, data }
    }

    /// Scalar multiplication.
    fn scale(&self, s: f64) -> Self {
        let data: Vec<f64> = self.data.iter().map(|x| x * s).collect();
        Mat { rows: self.rows, cols: self.cols, data }
    }

    /// Invert a square matrix via Gauss-Jordan elimination.
    fn inv(&self) -> Result<Self> {
        let n = self.rows;
        if n != self.cols {
            return Err(TimeSeriesError::ComputationError("Non-square matrix inversion".into()));
        }
        // Augmented [A | I]
        let mut aug: Vec<f64> = Vec::with_capacity(n * 2 * n);
        for i in 0..n {
            for j in 0..n { aug.push(self.get(i, j)); }
            for j in 0..n { aug.push(if i == j { 1.0 } else { 0.0 }); }
        }
        let cols = 2 * n;
        for col in 0..n {
            // Partial pivot
            let mut pivot_row = col;
            let mut max_val = aug[col * cols + col].abs();
            for row in (col + 1)..n {
                let v = aug[row * cols + col].abs();
                if v > max_val { max_val = v; pivot_row = row; }
            }
            if max_val < 1e-15 {
                return Err(TimeSeriesError::ComputationError("Singular matrix".into()));
            }
            if pivot_row != col {
                for j in 0..cols { aug.swap(col * cols + j, pivot_row * cols + j); }
            }
            let diag = aug[col * cols + col];
            for j in 0..cols { aug[col * cols + j] /= diag; }
            for row in 0..n {
                if row == col { continue; }
                let factor = aug[row * cols + col];
                for j in 0..cols { aug[row * cols + j] -= factor * aug[col * cols + j]; }
            }
        }
        let mut inv = Mat::zeros(n, n);
        for i in 0..n {
            for j in 0..n { inv.set(i, j, aug[i * cols + (n + j)]); }
        }
        Ok(inv)
    }

    /// Extract a column as a Vec.
    fn col_vec(&self, c: usize) -> Vec<f64> {
        (0..self.rows).map(|r| self.get(r, c)).collect()
    }

    /// From a column vector (Nx1).
    fn from_col(v: &[f64]) -> Self {
        Mat { rows: v.len(), cols: 1, data: v.to_vec() }
    }
}

// ---------------------------------------------------------------------------
// Kalman filter/smoother core (internal, works on Mat)
// ---------------------------------------------------------------------------

struct KfOutput {
    /// filtered state means: T × n_state
    filtered_mean: Vec<Vec<f64>>,
    /// filtered state covariances: T × (n_state × n_state flat)
    filtered_cov: Vec<Mat>,
    /// predicted state means: T × n_state
    predicted_mean: Vec<Vec<f64>>,
    /// predicted state covariances
    predicted_cov: Vec<Mat>,
    /// innovations (scalar observations)
    innovations: Vec<f64>,
    /// innovation variances
    innov_var: Vec<f64>,
    /// total log-likelihood
    log_lik: f64,
}

/// Run a univariate Kalman filter.
///
/// # Arguments
/// * `y`   – scalar observations (length T)
/// * `f`   – state transition matrix (n × n)
/// * `h`   – observation matrix (1 × n)
/// * `q`   – process noise covariance (n × n)
/// * `r`   – observation noise variance (scalar, stored as 1×1)
/// * `m0`  – initial state mean (n)
/// * `p0`  – initial state covariance (n × n)
fn kalman_filter_internal(
    y: &[f64],
    f: &Mat, h: &Mat, q: &Mat, r: f64,
    m0: &[f64], p0: &Mat,
) -> KfOutput {
    let n = f.rows;
    let t = y.len();
    let mut filt_mean = Vec::with_capacity(t);
    let mut filt_cov  = Vec::with_capacity(t);
    let mut pred_mean = Vec::with_capacity(t);
    let mut pred_cov  = Vec::with_capacity(t);
    let mut innovations = Vec::with_capacity(t);
    let mut innov_var   = Vec::with_capacity(t);

    let mut m = m0.to_vec();
    let mut p = p0.clone();
    let mut log_lik = 0.0f64;

    let f_t = f.t();
    let h_t = h.t();

    for obs in y.iter() {
        // --- Predict ---
        // m_pred = F * m
        let m_mat = Mat::from_col(&m);
        let mp_mat = f.mul(&m_mat);
        let m_pred: Vec<f64> = mp_mat.col_vec(0);
        // P_pred = F P F' + Q
        let fpt = f.mul(&p).mul(&f_t).add(q);

        pred_mean.push(m_pred.clone());
        pred_cov.push(fpt.clone());

        // --- Update ---
        // v = y_t - H m_pred  (scalar)
        let h_m: f64 = (0..n).map(|j| h.get(0, j) * m_pred[j]).sum();
        let v = obs - h_m;
        // S = H P_pred H' + R  (scalar)
        let hp = h.mul(&fpt);
        let s: f64 = (0..n).map(|j| hp.get(0, j) * h.get(0, j)).sum::<f64>() + r;
        let s_safe = if s < 1e-15 { 1e-15 } else { s };
        // K = P_pred H' / S  (n×1)
        let k_mat = fpt.mul(&h_t).scale(1.0 / s_safe);
        // m_filt = m_pred + K v
        let m_filt: Vec<f64> = (0..n).map(|i| m_pred[i] + k_mat.get(i, 0) * v).collect();
        // P_filt = (I - K H) P_pred
        let mut kh = Mat::zeros(n, n);
        for i in 0..n { for j in 0..n { kh.set(i, j, k_mat.get(i, 0) * h.get(0, j)); } }
        let ikm = Mat::eye(n).sub(&kh);
        let p_filt = ikm.mul(&fpt);

        // log-likelihood contribution
        log_lik -= 0.5 * ((2.0 * PI * s_safe).ln() + v * v / s_safe);

        filt_mean.push(m_filt.clone());
        filt_cov.push(p_filt.clone());
        innovations.push(v);
        innov_var.push(s_safe);

        m = m_filt;
        p = p_filt;
    }

    KfOutput {
        filtered_mean: filt_mean,
        filtered_cov: filt_cov,
        predicted_mean: pred_mean,
        predicted_cov: pred_cov,
        innovations,
        innov_var,
        log_lik,
    }
}

/// RTS backward smoother.
fn rts_smoother_internal(
    kf: &KfOutput,
    f: &Mat,
) -> (Vec<Vec<f64>>, Vec<Mat>) {
    let t = kf.filtered_mean.len();
    if t == 0 { return (vec![], vec![]); }
    let n = f.rows;
    let f_t = f.t();

    let mut smoothed_mean: Vec<Vec<f64>> = kf.filtered_mean.clone();
    let mut smoothed_cov: Vec<Mat>       = kf.filtered_cov.clone();

    for t_idx in (0..t - 1).rev() {
        let p_filt = &kf.filtered_cov[t_idx];
        let p_pred = &kf.predicted_cov[t_idx + 1];
        // Gain: G = P_filt F' P_pred^{-1}
        let p_pred_inv = match p_pred.inv() {
            Ok(inv) => inv,
            Err(_)  => {
                // Fallback: use a large-variance regularised inverse
                let mut reg = p_pred.clone();
                for i in 0..n { reg.add_assign(i, i, 1e-8); }
                match reg.inv() {
                    Ok(inv) => inv,
                    Err(_)  => Mat::eye(n).scale(1.0 / 1e-8),
                }
            }
        };
        let g = p_filt.mul(&f_t).mul(&p_pred_inv);

        let sm_next  = &smoothed_mean[t_idx + 1];
        let pm_next  = &kf.predicted_mean[t_idx + 1];
        let sp_next  = &smoothed_cov[t_idx + 1];
        let pm_filt  = &kf.filtered_mean[t_idx];

        // m_smooth = m_filt + G (m_smooth_{t+1} - m_pred_{t+1})
        let diff: Vec<f64> = (0..n).map(|i| sm_next[i] - pm_next[i]).collect();
        let g_diff = g.mul(&Mat::from_col(&diff));
        let m_smooth: Vec<f64> = (0..n).map(|i| pm_filt[i] + g_diff.get(i, 0)).collect();

        // P_smooth = P_filt + G (P_smooth_{t+1} - P_pred_{t+1}) G'
        let dp = sp_next.sub(p_pred);
        let g_t = g.t();
        let p_smooth = p_filt.add(&g.mul(&dp).mul(&g_t));

        smoothed_mean[t_idx] = m_smooth;
        smoothed_cov[t_idx]  = p_smooth;
    }

    (smoothed_mean, smoothed_cov)
}

// ---------------------------------------------------------------------------
// Nelder-Mead optimiser (unbounded, minimises a scalar closure)
// ---------------------------------------------------------------------------

/// Nelder-Mead simplex optimiser. Minimises `f(x)`.
fn nelder_mead<F>(
    f: F,
    x0: &[f64],
    max_iter: usize,
    tol: f64,
) -> Vec<f64>
where
    F: Fn(&[f64]) -> f64,
{
    let n = x0.len();
    let mut simplex: Vec<Vec<f64>> = Vec::with_capacity(n + 1);
    simplex.push(x0.to_vec());
    for i in 0..n {
        let mut v = x0.to_vec();
        v[i] += if v[i].abs() < 1e-10 { 0.1 } else { 0.05 * v[i].abs() + 0.05 };
        simplex.push(v);
    }
    let mut vals: Vec<f64> = simplex.iter().map(|v| f(v)).collect();

    for _ in 0..max_iter {
        // Sort by function value
        let mut idx: Vec<usize> = (0..n + 1).collect();
        idx.sort_by(|&a, &b| vals[a].partial_cmp(&vals[b]).unwrap_or(std::cmp::Ordering::Equal));
        let simplex_sorted: Vec<Vec<f64>> = idx.iter().map(|&i| simplex[i].clone()).collect();
        let vals_sorted: Vec<f64>         = idx.iter().map(|&i| vals[i]).collect();
        simplex = simplex_sorted;
        vals    = vals_sorted;

        // Convergence check
        let span = vals.last().copied().unwrap_or(f64::INFINITY) - vals[0];
        if span < tol { break; }

        // Centroid of all but worst
        let mut centroid = vec![0.0f64; n];
        for v in &simplex[..n] { for (j, &x) in v.iter().enumerate() { centroid[j] += x / n as f64; } }

        // Reflect
        let reflect: Vec<f64> = (0..n).map(|j| centroid[j] + 1.0 * (centroid[j] - simplex[n][j])).collect();
        let fr = f(&reflect);
        if fr < vals[0] {
            // Expand
            let expand: Vec<f64> = (0..n).map(|j| centroid[j] + 2.0 * (reflect[j] - centroid[j])).collect();
            let fe = f(&expand);
            if fe < fr { simplex[n] = expand; vals[n] = fe; }
            else        { simplex[n] = reflect; vals[n] = fr; }
        } else if fr < vals[n - 1] {
            simplex[n] = reflect; vals[n] = fr;
        } else {
            // Contract
            let contract: Vec<f64> = (0..n).map(|j| centroid[j] + 0.5 * (simplex[n][j] - centroid[j])).collect();
            let fc = f(&contract);
            if fc < vals[n] {
                simplex[n] = contract; vals[n] = fc;
            } else {
                // Shrink
                for i in 1..=n {
                    let shrunk: Vec<f64> = (0..n).map(|j| simplex[0][j] + 0.5 * (simplex[i][j] - simplex[0][j])).collect();
                    vals[i]    = f(&shrunk);
                    simplex[i] = shrunk;
                }
            }
        }
    }
    simplex[0].clone()
}

// Convert log-parameterised variances to natural (ensures positivity).
#[inline]
fn from_log_params(p: &[f64]) -> Vec<f64> { p.iter().map(|x| x.exp()).collect() }

// ---------------------------------------------------------------------------
// Public types
// ---------------------------------------------------------------------------

/// Output from Kalman filter / smoother.
#[derive(Clone, Debug)]
pub struct KalmanState {
    /// Filtered (or smoothed) state means, one per time point.
    pub state_means: Vec<Vec<f64>>,
    /// Filtered (or smoothed) state covariances, one per time point.
    pub state_covs: Vec<Vec<Vec<f64>>>,
    /// Prediction innovations (one-step-ahead residuals).
    pub innovations: Vec<f64>,
    /// Innovation variances.
    pub innov_vars: Vec<f64>,
    /// Total log-likelihood of the observations.
    pub log_likelihood: f64,
}

impl KalmanState {
    fn from_filter(kf: &KfOutput) -> Self {
        let state_covs = kf.filtered_cov.iter()
            .map(|m| mat_to_nested(m))
            .collect();
        KalmanState {
            state_means: kf.filtered_mean.clone(),
            state_covs,
            innovations: kf.innovations.clone(),
            innov_vars:  kf.innov_var.clone(),
            log_likelihood: kf.log_lik,
        }
    }
    fn from_smoother(sm: (Vec<Vec<f64>>, Vec<Mat>), log_lik: f64, innov: &[f64], iv: &[f64]) -> Self {
        let state_covs = sm.1.iter().map(|m| mat_to_nested(m)).collect();
        KalmanState {
            state_means: sm.0,
            state_covs,
            innovations: innov.to_vec(),
            innov_vars:  iv.to_vec(),
            log_likelihood: log_lik,
        }
    }
}

fn mat_to_nested(m: &Mat) -> Vec<Vec<f64>> {
    (0..m.rows).map(|r| (0..m.cols).map(|c| m.get(r, c)).collect()).collect()
}

// ---------------------------------------------------------------------------
// 1. Local Level Model  y_t = μ_t + ε_t,  μ_{t+1} = μ_t + η_t
// ---------------------------------------------------------------------------

/// Harvey's Local Level Model (random walk plus noise).
///
/// State equation:  `μ_{t+1} = μ_t + η_t`,  η_t ~ N(0, σ²_η)
/// Observation:     `y_t = μ_t + ε_t`,       ε_t ~ N(0, σ²_ε)
#[derive(Clone, Debug)]
pub struct LocalLevel {
    /// Process (level) noise variance σ²_η.
    pub level_var: f64,
    /// Observation noise variance σ²_ε.
    pub obs_var: f64,
}

impl LocalLevel {
    /// Fit a Local Level model by MLE using prediction-error decomposition.
    pub fn fit(data: &[f64]) -> Result<Self> {
        if data.len() < 3 {
            return Err(TimeSeriesError::InsufficientData {
                message: "LocalLevel::fit requires at least 3 observations".into(),
                required: 3,
                actual: data.len(),
            });
        }
        let neg_ll = |p: &[f64]| -> f64 {
            let vars = from_log_params(p);
            let (lv, ov) = (vars[0], vars[1]);
            let model = LocalLevel { level_var: lv, obs_var: ov };
            let kf = model.build_kf(data);
            -kf.log_lik
        };
        // Initialise from empirical variance
        let mean = data.iter().sum::<f64>() / data.len() as f64;
        let var  = data.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / data.len() as f64;
        let x0   = vec![(var * 0.3 + 1e-6).ln(), (var * 0.7 + 1e-6).ln()];
        let best = nelder_mead(neg_ll, &x0, 2000, 1e-8);
        let vars = from_log_params(&best);
        Ok(LocalLevel { level_var: vars[0].max(1e-10), obs_var: vars[1].max(1e-10) })
    }

    /// Run the Kalman filter and return filtered states.
    pub fn filter(&self, data: &[f64]) -> Result<KalmanState> {
        if data.is_empty() {
            return Err(TimeSeriesError::InsufficientData {
                message: "filter requires at least one observation".into(),
                required: 1,
                actual: 0,
            });
        }
        let kf = self.build_kf(data);
        Ok(KalmanState::from_filter(&kf))
    }

    /// Run the RTS smoother and return smoothed states.
    pub fn smoother(&self, data: &[f64]) -> Result<KalmanState> {
        if data.is_empty() {
            return Err(TimeSeriesError::InsufficientData {
                message: "smoother requires at least one observation".into(),
                required: 1,
                actual: 0,
            });
        }
        let kf = self.build_kf(data);
        let f  = Self::make_f();
        let sm = rts_smoother_internal(&kf, &f);
        Ok(KalmanState::from_smoother(sm, kf.log_lik, &kf.innovations, &kf.innov_var))
    }

    fn make_f() -> Mat { Mat::eye(1) }

    fn build_kf(&self, data: &[f64]) -> KfOutput {
        let f  = Self::make_f();
        let h  = Mat::eye(1);
        let mut q = Mat::zeros(1, 1); q.set(0, 0, self.level_var);
        let r  = self.obs_var;
        let m0 = vec![data[0]];
        let mut p0 = Mat::zeros(1, 1); p0.set(0, 0, 1e6);
        kalman_filter_internal(data, &f, &h, &q, r, &m0, &p0)
    }
}

// ---------------------------------------------------------------------------
// 2. Local Linear Trend  (level + slope)
// ---------------------------------------------------------------------------

/// Harvey's Local Linear Trend model.
///
/// Level: `μ_{t+1} = μ_t + β_t + η_t`,  η ~ N(0, σ²_η)
/// Slope: `β_{t+1} = β_t + ζ_t`,        ζ ~ N(0, σ²_ζ)
/// Obs:   `y_t = μ_t + ε_t`,             ε ~ N(0, σ²_ε)
#[derive(Clone, Debug)]
pub struct LocalLinearTrend {
    /// Level noise variance σ²_η.
    pub level_var: f64,
    /// Slope noise variance σ²_ζ.
    pub slope_var: f64,
    /// Observation noise variance σ²_ε.
    pub obs_var: f64,
}

impl LocalLinearTrend {
    /// Fit by MLE.
    pub fn fit(data: &[f64]) -> Result<Self> {
        if data.len() < 4 {
            return Err(TimeSeriesError::InsufficientData {
                message: "LocalLinearTrend::fit requires at least 4 observations".into(),
                required: 4,
                actual: data.len(),
            });
        }
        let neg_ll = |p: &[f64]| -> f64 {
            let vars = from_log_params(p);
            let model = LocalLinearTrend {
                level_var: vars[0], slope_var: vars[1], obs_var: vars[2],
            };
            let kf = model.build_kf(data);
            -kf.log_lik
        };
        let mean = data.iter().sum::<f64>() / data.len() as f64;
        let var  = data.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / data.len() as f64;
        let x0   = vec![(var * 0.2 + 1e-6).ln(), (var * 0.1 + 1e-6).ln(), (var * 0.7 + 1e-6).ln()];
        let best = nelder_mead(neg_ll, &x0, 3000, 1e-8);
        let vars = from_log_params(&best);
        Ok(LocalLinearTrend {
            level_var: vars[0].max(1e-10),
            slope_var: vars[1].max(1e-10),
            obs_var:   vars[2].max(1e-10),
        })
    }

    /// Run Kalman filter.
    pub fn filter(&self, data: &[f64]) -> Result<KalmanState> {
        if data.is_empty() {
            return Err(TimeSeriesError::InsufficientData {
                message: "filter requires at least one observation".into(),
                required: 1,
                actual: 0,
            });
        }
        let kf = self.build_kf(data);
        Ok(KalmanState::from_filter(&kf))
    }

    /// Run RTS smoother.
    pub fn smoother(&self, data: &[f64]) -> Result<KalmanState> {
        if data.is_empty() {
            return Err(TimeSeriesError::InsufficientData {
                message: "smoother requires at least one observation".into(),
                required: 1,
                actual: 0,
            });
        }
        let kf = self.build_kf(data);
        let (f, _, _, _) = Self::build_matrices(self.level_var, self.slope_var);
        let sm = rts_smoother_internal(&kf, &f);
        Ok(KalmanState::from_smoother(sm, kf.log_lik, &kf.innovations, &kf.innov_var))
    }

    fn build_matrices(level_var: f64, slope_var: f64) -> (Mat, Mat, Mat, Mat) {
        // State: [μ, β]'
        // F = [[1, 1], [0, 1]]
        let mut f = Mat::eye(2);
        f.set(0, 1, 1.0);
        // H = [1, 0]
        let mut h = Mat::zeros(1, 2);
        h.set(0, 0, 1.0);
        // Q = diag(level_var, slope_var)
        let mut q = Mat::zeros(2, 2);
        q.set(0, 0, level_var);
        q.set(1, 1, slope_var);
        // P0 = large diagonal
        let mut p0 = Mat::zeros(2, 2);
        p0.set(0, 0, 1e6);
        p0.set(1, 1, 1e6);
        (f, h, q, p0)
    }

    fn build_kf(&self, data: &[f64]) -> KfOutput {
        let (f, h, q, p0) = Self::build_matrices(self.level_var, self.slope_var);
        let r  = self.obs_var;
        let m0 = vec![data[0], 0.0];
        kalman_filter_internal(data, &f, &h, &q, r, &m0, &p0)
    }
}

// ---------------------------------------------------------------------------
// 3. Basic Structural Model (BSM)
// ---------------------------------------------------------------------------

/// Structural decomposition output.
#[derive(Clone, Debug)]
pub struct StructuralComponents {
    /// Smoothed trend component (μ_t).
    pub trend: Vec<f64>,
    /// Smoothed seasonal component (γ_t).
    pub seasonal: Vec<f64>,
    /// Irregular (residual) component (ε_t = y_t − trend − seasonal).
    pub irregular: Vec<f64>,
}

/// Forecast result.
#[derive(Clone, Debug)]
pub struct ForecastResult {
    /// Point forecasts.
    pub mean: Vec<f64>,
    /// Lower bound (95 % prediction interval).
    pub lower: Vec<f64>,
    /// Upper bound (95 % prediction interval).
    pub upper: Vec<f64>,
}

/// Harvey's Basic Structural Model with trigonometric seasonality.
///
/// ```text
/// y_t = μ_t + γ_t + ε_t
/// μ_{t+1} = μ_t + β_t + η_t     (level)
/// β_{t+1} = β_t + ζ_t            (slope)
/// γ_t = Σ_{j=1}^{⌊s/2⌋} γ_{j,t}  (trigonometric seasonal)
/// ```
///
/// Each seasonal harmonic j evolves as:
/// ```text
/// [γ_{j,t+1}]   [cos λ_j   sin λ_j] [γ_{j,t}]   [ω_{j,t}  ]
/// [γ*_{j,t+1}] = [-sin λ_j  cos λ_j] [γ*_{j,t}] + [ω*_{j,t} ]
/// ```
/// where λ_j = 2πj/s.
#[derive(Clone, Debug)]
pub struct BasicStructural {
    /// Level noise variance σ²_η.
    pub level_var: f64,
    /// Slope noise variance σ²_ζ.
    pub slope_var: f64,
    /// Seasonal disturbance variance σ²_ω (shared across harmonics).
    pub seasonal_var: f64,
    /// Observation noise variance σ²_ε.
    pub obs_var: f64,
    /// Seasonal period s.
    pub period: usize,
}

impl BasicStructural {
    /// Fit a BSM to `data` with given `period` by MLE.
    ///
    /// # Arguments
    /// * `data`   – univariate time series
    /// * `period` – seasonal period (e.g. 12 for monthly, 4 for quarterly)
    pub fn fit(data: &[f64], period: usize) -> Result<Self> {
        if period < 2 {
            return Err(TimeSeriesError::InvalidParameter {
                name: "period".into(),
                message: "period must be ≥ 2".into(),
            });
        }
        if data.len() < 2 * period {
            return Err(TimeSeriesError::InsufficientData {
                message: format!("BSM fit requires at least {} observations", 2 * period),
                required: 2 * period,
                actual: data.len(),
            });
        }
        let neg_ll = |p: &[f64]| -> f64 {
            let vars = from_log_params(p);
            let model = BasicStructural {
                level_var:    vars[0],
                slope_var:    vars[1],
                seasonal_var: vars[2],
                obs_var:      vars[3],
                period,
            };
            let (f, h, q, p0, m0) = model.build_matrices(data);
            let kf = kalman_filter_internal(data, &f, &h, &q, vars[3], &m0, &p0);
            -kf.log_lik
        };
        let mean = data.iter().sum::<f64>() / data.len() as f64;
        let var  = data.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / data.len() as f64;
        let x0 = vec![
            (var * 0.15 + 1e-6).ln(),
            (var * 0.05 + 1e-6).ln(),
            (var * 0.10 + 1e-6).ln(),
            (var * 0.70 + 1e-6).ln(),
        ];
        let best = nelder_mead(neg_ll, &x0, 4000, 1e-8);
        let vars = from_log_params(&best);
        Ok(BasicStructural {
            level_var:    vars[0].max(1e-10),
            slope_var:    vars[1].max(1e-10),
            seasonal_var: vars[2].max(1e-10),
            obs_var:      vars[3].max(1e-10),
            period,
        })
    }

    /// Number of seasonal harmonics: ⌊period/2⌋.
    pub fn n_harmonics(&self) -> usize { self.period / 2 }

    /// State dimension: 2 (trend/slope) + 2 * n_harmonics.
    pub fn n_states(&self) -> usize { 2 + 2 * self.n_harmonics() }

    /// Build the state-space matrices for this model.
    ///
    /// Returns (F, H, Q, P0, m0).
    fn build_matrices(&self, data: &[f64]) -> (Mat, Mat, Mat, Mat, Vec<f64>) {
        let ns = self.n_states();
        let nh = self.n_harmonics();

        // --- Transition matrix F ---
        // [1  1  |  0  0  ...  ]    trend-level row
        // [0  1  |  0  0  ...  ]    trend-slope row
        // [       | Rj blocks  ]    seasonal harmonic blocks
        let mut f = Mat::zeros(ns, ns);
        // Trend block
        f.set(0, 0, 1.0);
        f.set(0, 1, 1.0);
        f.set(1, 1, 1.0);
        // Seasonal harmonic blocks (each 2×2 rotation)
        for j in 1..=nh {
            let lam = 2.0 * PI * j as f64 / self.period as f64;
            let (c, s) = (lam.cos(), lam.sin());
            let base = 2 + (j - 1) * 2;
            f.set(base,     base,     c);
            f.set(base,     base + 1, s);
            f.set(base + 1, base,    -s);
            f.set(base + 1, base + 1, c);
        }

        // --- Observation matrix H: y_t = μ_t + Σ γ_{j,t} ---
        let mut h = Mat::zeros(1, ns);
        h.set(0, 0, 1.0);
        for j in 0..nh { h.set(0, 2 + j * 2, 1.0); }

        // --- Process noise Q ---
        let mut q = Mat::zeros(ns, ns);
        q.set(0, 0, self.level_var);
        q.set(1, 1, self.slope_var);
        for j in 0..nh {
            let base = 2 + j * 2;
            q.set(base,     base,     self.seasonal_var);
            q.set(base + 1, base + 1, self.seasonal_var);
        }

        // --- Initial covariance P0 (diffuse) ---
        let mut p0 = Mat::zeros(ns, ns);
        for i in 0..ns { p0.set(i, i, 1e6); }

        // --- Initial state mean m0 ---
        let first = if data.is_empty() { 0.0 } else { data[0] };
        let mut m0 = vec![0.0f64; ns];
        m0[0] = first;

        (f, h, q, p0, m0)
    }

    fn build_kf(&self, data: &[f64]) -> KfOutput {
        let (f, h, q, p0, m0) = self.build_matrices(data);
        kalman_filter_internal(data, &f, &h, &q, self.obs_var, &m0, &p0)
    }

    /// Kalman filter output.
    pub fn filter(&self, data: &[f64]) -> Result<KalmanState> {
        if data.is_empty() {
            return Err(TimeSeriesError::InsufficientData {
                message: "filter requires at least one observation".into(),
                required: 1,
                actual: 0,
            });
        }
        let kf = self.build_kf(data);
        Ok(KalmanState::from_filter(&kf))
    }

    /// RTS smoother output.
    pub fn smoother(&self, data: &[f64]) -> Result<KalmanState> {
        if data.is_empty() {
            return Err(TimeSeriesError::InsufficientData {
                message: "smoother requires at least one observation".into(),
                required: 1,
                actual: 0,
            });
        }
        let kf = self.build_kf(data);
        let (f, _, _, _, _) = self.build_matrices(data);
        let sm = rts_smoother_internal(&kf, &f);
        Ok(KalmanState::from_smoother(sm, kf.log_lik, &kf.innovations, &kf.innov_var))
    }
}

// ---------------------------------------------------------------------------
// 4. Structural decomposition
// ---------------------------------------------------------------------------

/// Decompose a time series into trend, seasonal, and irregular components
/// using the smoothed states from a fitted `BasicStructural` model.
///
/// # Arguments
/// * `model` – fitted BSM
/// * `data`  – original observations
///
/// Returns [`StructuralComponents`] containing smoothed trend, seasonal, and irregular.
pub fn decompose(model: &BasicStructural, data: &[f64]) -> Result<StructuralComponents> {
    if data.is_empty() {
        return Err(TimeSeriesError::InsufficientData {
            message: "decompose requires at least one observation".into(),
            required: 1,
            actual: 0,
        });
    }
    let smoothed = model.smoother(data)?;
    let nh = model.n_harmonics();
    let n  = data.len();

    let mut trend    = Vec::with_capacity(n);
    let mut seasonal = Vec::with_capacity(n);

    for t in 0..n {
        let st = &smoothed.state_means[t];
        trend.push(st[0]); // μ_t is state index 0
        // Seasonal = sum of first element of each harmonic pair
        let s: f64 = (0..nh).map(|j| st[2 + j * 2]).sum();
        seasonal.push(s);
    }

    let irregular: Vec<f64> = (0..n)
        .map(|t| data[t] - trend[t] - seasonal[t])
        .collect();

    Ok(StructuralComponents { trend, seasonal, irregular })
}

// ---------------------------------------------------------------------------
// 5. Forecasting
// ---------------------------------------------------------------------------

/// Forecast `h` steps ahead from a fitted `BasicStructural` model.
///
/// Uses the Kalman-filtered last state as the initial distribution and
/// propagates uncertainty forward through the state equations.
///
/// # Arguments
/// * `model` – fitted BSM
/// * `data`  – observed history (used to run Kalman filter)
/// * `h`     – forecast horizon (number of steps)
pub fn forecast(model: &BasicStructural, data: &[f64], h: usize) -> Result<ForecastResult> {
    if data.is_empty() {
        return Err(TimeSeriesError::InsufficientData {
            message: "forecast requires at least one observation".into(),
            required: 1,
            actual: 0,
        });
    }
    if h == 0 {
        return Ok(ForecastResult { mean: vec![], lower: vec![], upper: vec![] });
    }

    let (f, h_mat, q, _, _) = model.build_matrices(data);
    let f_t = f.t();
    let kf  = model.build_kf(data);

    let t_last = kf.filtered_mean.len() - 1;
    let mut m = kf.filtered_mean[t_last].clone();
    let mut p = kf.filtered_cov[t_last].clone();

    let mut means  = Vec::with_capacity(h);
    let mut lowers = Vec::with_capacity(h);
    let mut uppers = Vec::with_capacity(h);
    let z95 = 1.959_964f64; // 95 % CI

    for _ in 0..h {
        // Predict one step
        let m_mat  = Mat::from_col(&m);
        let mp_mat = f.mul(&m_mat);
        m = mp_mat.col_vec(0);
        p = f.mul(&p).mul(&f_t).add(&q);

        // Forecast mean: y_pred = H m
        let n = f.rows;
        let y_pred: f64 = (0..n).map(|j| h_mat.get(0, j) * m[j]).sum();

        // Forecast variance: S = H P H' + R
        let hp = h_mat.mul(&p);
        let s: f64 = (0..n).map(|j| hp.get(0, j) * h_mat.get(0, j)).sum::<f64>() + model.obs_var;
        let std_dev = s.sqrt();

        means.push(y_pred);
        lowers.push(y_pred - z95 * std_dev);
        uppers.push(y_pred + z95 * std_dev);
    }

    Ok(ForecastResult { mean: means, lower: lowers, upper: uppers })
}

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

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

    fn linspace(n: usize, slope: f64, noise_amp: f64) -> Vec<f64> {
        (0..n).map(|i| slope * i as f64 + noise_amp * (i as f64 * 0.7).sin()).collect()
    }

    fn seasonal_series(n: usize, period: usize, amp: f64, noise: f64) -> Vec<f64> {
        (0..n).map(|i| {
            let s = amp * (2.0 * PI * i as f64 / period as f64).sin();
            let trend = 0.05 * i as f64;
            s + trend + noise * (i as f64 * 1.3).cos()
        }).collect()
    }

    // --- LocalLevel tests ---

    #[test]
    fn test_local_level_fit_smoke() {
        let data: Vec<f64> = (0..30).map(|i| 5.0 + (i as f64 * 0.2).sin() * 0.5).collect();
        let model = LocalLevel::fit(&data).expect("fit ok");
        assert!(model.level_var > 0.0, "level_var must be positive");
        assert!(model.obs_var > 0.0, "obs_var must be positive");
    }

    #[test]
    fn test_local_level_filter_length() {
        let data: Vec<f64> = (0..20).map(|i| 3.0 + (i as f64 * 0.3).cos()).collect();
        let model = LocalLevel { level_var: 0.1, obs_var: 0.5 };
        let ks = model.filter(&data).expect("filter ok");
        assert_eq!(ks.state_means.len(), 20);
        assert_eq!(ks.innovations.len(), 20);
    }

    #[test]
    fn test_local_level_smoother_length() {
        let data: Vec<f64> = (0..15).map(|i| (i as f64 * 0.4).sin()).collect();
        let model = LocalLevel { level_var: 0.2, obs_var: 0.3 };
        let ks = model.smoother(&data).expect("smoother ok");
        assert_eq!(ks.state_means.len(), 15);
    }

    #[test]
    fn test_local_level_filter_empty_err() {
        let model = LocalLevel { level_var: 0.1, obs_var: 0.1 };
        assert!(model.filter(&[]).is_err());
    }

    #[test]
    fn test_local_level_constant_series() {
        // For a constant series, the filtered level should converge to the constant.
        let data = vec![5.0f64; 50];
        let model = LocalLevel { level_var: 0.01, obs_var: 0.1 };
        let ks = model.filter(&data).expect("filter ok");
        let last_level = ks.state_means.last().and_then(|v| v.first()).copied().unwrap_or(0.0);
        assert!((last_level - 5.0).abs() < 0.5, "last level={last_level}");
    }

    #[test]
    fn test_local_level_log_likelihood_finite() {
        let data: Vec<f64> = (0..20).map(|i| (i as f64 * 0.5).sin() * 2.0).collect();
        let model = LocalLevel { level_var: 0.1, obs_var: 0.5 };
        let ks = model.filter(&data).expect("filter ok");
        assert!(ks.log_likelihood.is_finite());
        assert!(ks.log_likelihood < 0.0, "ll should be negative");
    }

    // --- LocalLinearTrend tests ---

    #[test]
    fn test_llt_fit_smoke() {
        let data = linspace(40, 0.5, 0.3);
        let model = LocalLinearTrend::fit(&data).expect("fit ok");
        assert!(model.level_var > 0.0);
        assert!(model.slope_var > 0.0);
        assert!(model.obs_var > 0.0);
    }

    #[test]
    fn test_llt_filter_tracks_trend() {
        let data: Vec<f64> = (0..25).map(|i| i as f64 + 0.1 * (i as f64).cos()).collect();
        let model = LocalLinearTrend { level_var: 0.01, slope_var: 0.001, obs_var: 0.5 };
        let ks = model.filter(&data).expect("filter ok");
        let last_level = ks.state_means.last().and_then(|v| v.first()).copied().unwrap_or(0.0);
        assert!((last_level - 24.0).abs() < 5.0, "last_level={last_level}");
    }

    #[test]
    fn test_llt_smoother_length() {
        let data = linspace(30, 1.0, 0.2);
        let model = LocalLinearTrend { level_var: 0.05, slope_var: 0.005, obs_var: 0.3 };
        let ks = model.smoother(&data).expect("smoother ok");
        assert_eq!(ks.state_means.len(), 30);
        assert_eq!(ks.state_means[0].len(), 2, "state dim should be 2");
    }

    // --- BasicStructural tests ---

    #[test]
    fn test_bsm_fit_smoke() {
        let data = seasonal_series(48, 12, 2.0, 0.3);
        let model = BasicStructural::fit(&data, 12).expect("BSM fit ok");
        assert!(model.level_var > 0.0);
        assert!(model.slope_var > 0.0);
        assert!(model.seasonal_var > 0.0);
        assert!(model.obs_var > 0.0);
        assert_eq!(model.period, 12);
    }

    #[test]
    fn test_bsm_invalid_period() {
        let data = vec![1.0f64; 20];
        assert!(BasicStructural::fit(&data, 1).is_err());
    }

    #[test]
    fn test_bsm_insufficient_data() {
        let data = vec![1.0f64; 5];
        assert!(BasicStructural::fit(&data, 12).is_err());
    }

    #[test]
    fn test_bsm_n_states_quarterly() {
        let model = BasicStructural {
            level_var: 0.1, slope_var: 0.01, seasonal_var: 0.05, obs_var: 0.3, period: 4,
        };
        // period=4 → harmonics=2 → states = 2 + 2*2 = 6
        assert_eq!(model.n_harmonics(), 2);
        assert_eq!(model.n_states(), 6);
    }

    #[test]
    fn test_bsm_filter_length() {
        let data = seasonal_series(24, 4, 1.5, 0.2);
        let model = BasicStructural {
            level_var: 0.1, slope_var: 0.01, seasonal_var: 0.1, obs_var: 0.4, period: 4,
        };
        let ks = model.filter(&data).expect("filter ok");
        assert_eq!(ks.state_means.len(), 24);
        assert_eq!(ks.state_means[0].len(), model.n_states());
    }

    #[test]
    fn test_bsm_smoother_length() {
        let data = seasonal_series(20, 4, 1.0, 0.1);
        let model = BasicStructural {
            level_var: 0.1, slope_var: 0.01, seasonal_var: 0.1, obs_var: 0.3, period: 4,
        };
        let ks = model.smoother(&data).expect("smoother ok");
        assert_eq!(ks.state_means.len(), 20);
    }

    // --- Decomposition tests ---

    #[test]
    fn test_decompose_lengths() {
        let data = seasonal_series(24, 4, 2.0, 0.2);
        let model = BasicStructural {
            level_var: 0.1, slope_var: 0.01, seasonal_var: 0.1, obs_var: 0.3, period: 4,
        };
        let comps = decompose(&model, &data).expect("decompose ok");
        assert_eq!(comps.trend.len(),    24);
        assert_eq!(comps.seasonal.len(), 24);
        assert_eq!(comps.irregular.len(), 24);
    }

    #[test]
    fn test_decompose_reconstruction() {
        // trend + seasonal + irregular should approximate original series.
        let data = seasonal_series(24, 4, 2.0, 0.2);
        let model = BasicStructural {
            level_var: 0.1, slope_var: 0.01, seasonal_var: 0.1, obs_var: 0.3, period: 4,
        };
        let comps = decompose(&model, &data).expect("decompose ok");
        for t in 0..24 {
            let recon = comps.trend[t] + comps.seasonal[t] + comps.irregular[t];
            assert!((recon - data[t]).abs() < 1e-9, "mismatch at t={t}: recon={recon} data={}", data[t]);
        }
    }

    #[test]
    fn test_decompose_values_finite() {
        let data = seasonal_series(24, 12, 3.0, 0.5);
        let model = BasicStructural {
            level_var: 0.2, slope_var: 0.02, seasonal_var: 0.2, obs_var: 0.5, period: 12,
        };
        let comps = decompose(&model, &data).expect("decompose ok");
        for t in 0..data.len() {
            assert!(comps.trend[t].is_finite(), "trend not finite at {t}");
            assert!(comps.seasonal[t].is_finite(), "seasonal not finite at {t}");
            assert!(comps.irregular[t].is_finite(), "irregular not finite at {t}");
        }
    }

    // --- Forecast tests ---

    #[test]
    fn test_forecast_length() {
        let data = seasonal_series(24, 4, 1.5, 0.2);
        let model = BasicStructural {
            level_var: 0.1, slope_var: 0.01, seasonal_var: 0.1, obs_var: 0.3, period: 4,
        };
        let fc = forecast(&model, &data, 8).expect("forecast ok");
        assert_eq!(fc.mean.len(),  8);
        assert_eq!(fc.lower.len(), 8);
        assert_eq!(fc.upper.len(), 8);
    }

    #[test]
    fn test_forecast_intervals_valid() {
        let data = seasonal_series(24, 4, 1.5, 0.2);
        let model = BasicStructural {
            level_var: 0.1, slope_var: 0.01, seasonal_var: 0.1, obs_var: 0.3, period: 4,
        };
        let fc = forecast(&model, &data, 6).expect("forecast ok");
        for i in 0..6 {
            assert!(fc.lower[i] <= fc.mean[i], "lower > mean at h={i}");
            assert!(fc.upper[i] >= fc.mean[i], "upper < mean at h={i}");
        }
    }

    #[test]
    fn test_forecast_intervals_widen() {
        // Prediction intervals should generally widen over the horizon.
        let data = seasonal_series(24, 4, 1.5, 0.2);
        let model = BasicStructural {
            level_var: 0.1, slope_var: 0.01, seasonal_var: 0.1, obs_var: 0.3, period: 4,
        };
        let fc = forecast(&model, &data, 8).expect("forecast ok");
        let w0 = fc.upper[0] - fc.lower[0];
        let w7 = fc.upper[7] - fc.lower[7];
        assert!(w7 >= w0, "interval did not widen: w0={w0} w7={w7}");
    }

    #[test]
    fn test_forecast_zero_horizon() {
        let data = seasonal_series(12, 4, 1.0, 0.1);
        let model = BasicStructural {
            level_var: 0.1, slope_var: 0.01, seasonal_var: 0.05, obs_var: 0.2, period: 4,
        };
        let fc = forecast(&model, &data, 0).expect("zero-h forecast ok");
        assert!(fc.mean.is_empty());
    }

    #[test]
    fn test_forecast_finite_values() {
        let data = seasonal_series(24, 4, 2.0, 0.3);
        let model = BasicStructural {
            level_var: 0.1, slope_var: 0.01, seasonal_var: 0.1, obs_var: 0.3, period: 4,
        };
        let fc = forecast(&model, &data, 12).expect("forecast ok");
        for i in 0..12 {
            assert!(fc.mean[i].is_finite(), "mean[{i}] is not finite");
            assert!(fc.lower[i].is_finite(), "lower[{i}] is not finite");
            assert!(fc.upper[i].is_finite(), "upper[{i}] is not finite");
        }
    }
}