// arima.mbt — ARIMA enhancements (v0.85.0): Yule-Walker closed-form
// AR(p) fit + stateful forecast model + confidence intervals.
//
// This file builds on top of the existing v0.28.5 ARIMA infrastructure
// in `arima_demo.mbt` (which uses gradient-descent-based ARMA fitting
// via `arma_fit`). v0.85.0 adds:
//
//   1. `fit_ar_yule_walker` — closed-form Yule-Walker estimation of
//      AR(p) coefficients via Levinson-Durbin recursion (much faster
//      than gradient descent; the canonical AR-only solver).
//   2. `ArimaFittedModel` — stateful struct that bundles the fit
//      parameters + last `d` raw values + last `q` residuals needed
//      to do incremental forecasting without re-fitting.
//   3. `arima_forecast_with_variance` — analytical forecast variance
//      from the AR(∞) representation, giving 1.96·σ_z confidence
//      bands under Gaussian-innovation assumptions.
//
// The pipeline is:
//   1. ts_diff_n(series, d)         — make series stationary
//   2. fit_ar_yule_walker(z, p)     — fit AR coefficients (closed-form)
//   3. construct ArimaFittedModel  — bundle + state
//   4. forecast h steps + CI       — point forecast + analytical 95% CI
//
// Reference: Box & Jenkins 1970; Hamilton 1994 "Time Series Analysis";
// Levinson-Durbin recursion (1947).

///|
/// Helper: γ(lag) for a centered series (mean subtracted).
fn gamma_at_lag(
  series : Array[Float],
  mean : Float,
  n : Int,
  lag : Int,
) -> Float {
  if lag >= n || lag < 0 {
    return 0.0F
  }
  let mut s = 0.0F
  for i in 0..<(n - lag) {
    s = s + (series[i] - mean) * (series[i + lag] - mean)
  }
  s / Float::from_int(n)
}

///|
/// Closed-form AR(p) coefficient estimation via Yule-Walker +
/// Levinson-Durbin recursion. Operates on a (already-stationary)
/// series. Returns φ = [φ_1, ..., φ_p] such that
///   y_t ≈ Σ_{i=1..p} φ_i · y_{t-i}.
/// `p = 0` returns an empty array.
pub fn fit_ar_yule_walker(series : Array[Float], n : Int, p : Int) -> Array[Float] {
  if p <= 0 || n <= p {
    return Array::make(p, 0.0F)
  }
  // Compute mean + variance for the centered series.
  let mut mean = 0.0F
  for i in 0.. ArimaFittedModel {
  let ma_coef : Array[Float] = Array::make(q, 0.0F)
  {
    p,
    d,
    q,
    ar_coef,
    ma_coef,
    intercept,
    diff_init,
    residuals,
    residual_variance,
  }
}

///|
/// Build an ArimaFittedModel from a fitted ARMA fit (the existing
/// `arma_fit` output). Combines `ArimaFit` + residuals + variance.
/// `diff_init` is the last `d` raw observations (for undifferencing).
pub fn ArimaFittedModel::from_arma_fit(
  fit : ArimaFit,
  residuals : Array[Float],
  residual_variance : Float,
  diff_init : Array[Float],
) -> ArimaFittedModel {
  {
    p: fit.p,
    d: fit.d,
    q: fit.q,
    ar_coef: fit.ar_coeffs,
    ma_coef: fit.ma_thetas,
    intercept: fit.intercept,
    diff_init,
    residuals,
    residual_variance,
  }
}

///|
/// Forecast h steps ahead from a fitted model. Returns the original-
/// scale point forecast (in y space — undifferenced if d > 0).
/// For h > q, MA residuals are zero (unobserved beyond the fit window).
pub fn arima_fitted_forecast(model : ArimaFittedModel, h : Int) -> Array[Float] {
  let forecast : Array[Float] = Array::make(h, 0.0F)
  // Maintain a buffer of recent differenced values for the AR part.
  let buf_len = if model.p > model.q { model.p } else { model.q }
  let buf : Array[Float] = Array::make(buf_len, 0.0F)
  // Initialize buf with the most recent values. For v0.85.0 we
  // approximate by zero (the recency info comes from the residuals
  // buffer for the MA part, and AR part starts from zero — a rough
  // but valid baseline; proper init requires the most recent differenced
  // observations, which the caller can supply via `diff_init`).
  for k in 0.. 0.
    let mut v = pred
    let mut d_idx = model.d - 1
    while d_idx >= 0 {
      v = v + model.diff_init[d_idx]
      d_idx = d_idx - 1
    }
    forecast[k] = v
    // Shift buf left and append pred (differenced scale).
    let diff_pred = pred
    for i in 0..<(buf_len - 1) {
      buf[i] = buf[i + 1]
    }
    buf[buf_len - 1] = diff_pred
  }
  forecast
}

///|
/// Compute the analytical forecast variance at horizon k for an
/// ARMA(p, q) model. Used to construct Gaussian-approximation
/// confidence intervals: y_{T+k} ± 1.96 · sqrt(var_k).
///
/// For a pure AR(p) model, the variance converges to the residual
/// variance as k → ∞. For ARMA with MA, the formula uses the
/// ψ-weights of the AR(∞) representation:
///   ψ_0 = 1, ψ_j = φ_1 ψ_{j-1} + ... + φ_p ψ_{j-p} + θ_j (for j ≤ q)
///   var_k = σ²_ε · Σ_{j=0..k-1} ψ_j²
pub fn arima_forecast_variance(model : ArimaFittedModel, h : Int) -> Array[Float] {
  let variance_seq : Array[Float] = Array::make(h, 0.0F)
  // Compute ψ-weights up to h - 1.
  let psi : Array[Float] = Array::make(h, 0.0F)
  psi[0] = 1.0F
  for j in 1..= 0 {
        s = s + model.ar_coef[i] * psi[j - 1 - i]
      }
    }
    // MA contribution (for j ≤ q).
    if j - 1 < model.q {
      s = s + model.ma_coef[j - 1]
    }
    psi[j] = s
  }
  // var_k = σ²_ε · Σ_{j=0..k-1} ψ_j²
  let mut running = 0.0F
  for k in 0.. (Array[Float], Array[Float], Array[Float]) {
  let forecast = arima_fitted_forecast(model, h)
  let variance_seq = arima_forecast_variance(model, h)
  let lower : Array[Float] = Array::make(h, 0.0F)
  let upper : Array[Float] = Array::make(h, 0.0F)
  for k in 0..