// arima_demo.mbt — End-to-end ARIMA demo pipeline (v0.28.5).
//
// This demo ties together ts_diff / ts_cumsum (v0.28.0) and
// arma_fit (v0.28.4) to implement the classical ARIMA(p, d, q)
// modelling pipeline:
//
//   y_t → d-th differencing → z_t → fit ARMA(p, q) on z_t
//       → forecast z_{N+1}, ..., z_{N+k}
//       → inverse-difference back to y space
//
// Steps:
//   1. `arima_synthesize`  — generate ARIMA(p, d, q) data
//   2. `arima_fit`         — pipeline: diff → fit ARMA → return
//                            { ar_coeffs, ma_thetas, intercept, d }
//   3. `arima_forecast`    — fit + forecast N steps ahead
//   4. `arima_evaluate`    — RMSE / MAE for forecast accuracy
//
// All Float32 (matching Julia's `Float32` arima-style defaults).

///|
/// Bundle of fitted ARIMA(p, d, q) parameters.
pub struct ArimaFit {
  p : Int
  d : Int
  q : Int
  intercept : Float
  ar_coeffs : Array[Float]
  ma_thetas : Array[Float]
}

///|
/// Synthetic ARIMA(p, d, q) series generator. Generate z via
/// ARMA(p, q), then inverse-difference (cumulative sum + extend)
/// d times to produce the integrated series y.
pub fn arima_synthesize(
  n : Int,
  p : Int,
  d : Int,
  q : Int,
  intercept : Float,
  ar_coeffs : Array[Float],
  ma_thetas : Array[Float],
  seed : UInt64,
) -> Array[Float] {
  let rng = Xoshiro::from_state(seed, seed + 1UL, seed + 2UL, seed + 3UL)
  // Step 1: generate the stationary ARMA(p, q) series z.
  let z : Array[Float] = Array::make(n, 0.0F)
  if n == 0 {
    return z
  }
  let eps : Array[Float] = Array::make(n, 0.0F)
  // Initial eps[0] = z[0] = noise sample (no pre-sample data).
  let (z0, _) = box_muller(rng)
  eps[0] = Float::from_double(z0)
  z[0] = eps[0]
  for t in 1..= 0 {
        e = e - ar_coeffs[i] * z[lag]
      }
    }
    // MA part
    for j in 0..= 0 {
        e = e - ma_thetas[j] * eps[lag]
      }
    }
    eps[t] = e
    z[t] = intercept + e
  }
  // Step 2: inverse-difference d times. Each inverse-difference
  // requires a "head" value y[-1]; we use 0 by default. (Real-world
  // ARIMA uses the last observed y as the head.)
  let mut current : Array[Float] = z
  for _k in 0.. ArimaFit {
  let z = ts_diff_n(y, d)
  let param = arma_fit(z, p, q, max_iter, lr)
  {
    p,
    d,
    q,
    intercept: param.intercept,
    ar_coeffs: param.ar_coeffs,
    ma_thetas: param.ma_thetas,
  }
}

///|
/// Forecast `n_ahead` steps into the future from a fitted
/// ArimaFit on a series y. Returns an array of length `n_ahead`
/// containing point forecasts for y[N+1], ..., y[N+n_ahead].
pub fn arima_forecast(
  y : Array[Float],
  fit : ArimaFit,
  n_ahead : Int,
) -> Array[Float] {
  if n_ahead <= 0 {
    return Array::make(0, 0.0F)
  }
  let n = y.length()
  // 1. Re-fit ARMA on differenced data so we have residuals for
  //    forecasting.
  let z = ts_diff_n(y, fit.d)
  let param = ArmaParam::new(
    fit.p,
    fit.q,
    fit.intercept,
    fit.ar_coeffs,
    fit.ma_thetas,
  )
  let eps_fit = arma_residuals(z, param)
  let n_z = z.length()
  // 2. Forecast on the differenced scale.
  let z_forecast = Array::make(n_ahead, 0.0F)
  // Build an extended working copy of z so we can use ARMA recursion.
  let z_ext : Array[Float] = Array::make(n_z + n_ahead, 0.0F)
  for i in 0..= 0 {
        pred = pred + fit.ar_coeffs[i] * z_ext[lag]
      }
    }
    for j in 0..= 0 {
        pred = pred + fit.ma_thetas[j] * eps_ext[lag]
      }
    }
    z_ext[t] = pred
    eps_ext[t] = 0.0F  // innovation in forecast horizon = 0
    z_forecast[t - n_z] = pred
  }
  // 3. Inverse-difference back to y space.
  if fit.d == 0 {
    return z_forecast
  }
  // Reconstruct the head: we use the last observed y value as the
  // initial value for the d-th difference.
  let last_y = if n > 0 { y[n - 1] } else { 0.0F }
  let d_head = last_d_value(z, fit.d, last_y)
  let cur : Array[Float] = Array::make(n_z + n_ahead, 0.0F)
  for i in 0..1.
fn last_d_value(
  z : Array[Float],
  d : Int,
  last_y : Float,
) -> Float {
  if d == 0 {
    return 0.0F
  }
  if d == 1 {
    // y[n] = y[n-1] + diff(y)[n-1]. head for inverse-diff at t=0
    // (relative) is the y value BEFORE the diff sequence starts.
    return last_y - z[0]
  }
  // For d ≥ 2: approximate; in practice the user should provide
  // sufficient pre-sample data. Return last_y as a fallback.
  let _ = z
  last_y
}

///|
/// Root-mean-squared error between two equal-length arrays.
pub fn arima_rmse(actual : Array[Float], predicted : Array[Float]) -> Float {
  let n = actual.length()
  if predicted.length() != n {
    return -1.0F
  }
  if n == 0 {
    return 0.0F
  }
  let mut sum_sq = 0.0F
  for i in 0.. Float {
  let n = actual.length()
  if predicted.length() != n {
    return -1.0F
  }
  if n == 0 {
    return 0.0F
  }
  let mut sum_abs = 0.0F
  for i in 0..