// RAID Metal kernels — FP32 with double-float ("compensated") accumulation.
// Apple GPUs have no FP64. Every accumulator is a pair (hi, lo) maintained with error-free transforms (TwoSum /
// Fast2Sum); partial results per (series, chunk) are written as float2 and combined in FP64 on the CPU in a fixed
// order, which makes results deterministic for a given chunk size and threadgroup width.
#include <metal_stdlib>
using namespace metal;

// ---- double-float arithmetic -----------------------------------------------------------------------------------
struct df { float hi; float lo; };
inline df df_zero() { return df{0.0f, 0.0f}; }
inline df two_sum(float a, float b) { float s = a + b; float bb = s - a; float e = (a - (s - bb)) + (b - bb); return df{s, e}; }
inline df df_add_f(df a, float b) { df s = two_sum(a.hi, b); s.lo += a.lo; float h = s.hi + s.lo; float l = s.lo - (h - s.hi); return df{h, l}; }
inline df df_add(df a, df b) { df s = two_sum(a.hi, b.hi); s.lo += a.lo + b.lo; float h = s.hi + s.lo; float l = s.lo - (h - s.hi); return df{h, l}; }
inline float2 df_pack(df a) { return float2(a.hi, a.lo); }

// Threadgroup tree reduction of double-float values (deterministic order).
inline df tg_reduce(df v, threadgroup df* sh, uint lid, uint n) {
  sh[lid] = v; threadgroup_barrier(mem_flags::mem_threadgroup);
  for (uint s = n / 2; s > 0; s >>= 1) { if (lid < s) sh[lid] = df_add(sh[lid], sh[lid + s]); threadgroup_barrier(mem_flags::mem_threadgroup); }
  return sh[0];
}

// ---- 1. fused realized-measure accumulators ---------------------------------------------------------------------
// One threadgroup per (segment, chunk). Inputs: x (float), offsets (int64 as two uints packed → we use uint offsets,
// the host guarantees < 2^32 elements), chunk descriptors (segment id, begin, end). Outputs per chunk: 5 float2:
// Σx², Σ|x_i||x_{i-1}| (pairs with i-1 inside the chunk; the host adds the boundary products), Σx²1{x<0}, Σx²1{x>0}, Σx⁴.
struct Chunk { uint seg; uint begin; uint end; uint pad; };
kernel void realized_fused(device const float* x [[buffer(0)]], device const Chunk* chunks [[buffer(1)]], device float2* out [[buffer(2)]],
                           uint gid [[threadgroup_position_in_grid]], uint lid [[thread_position_in_threadgroup]], uint tsz [[threads_per_threadgroup]]) {
  threadgroup df sh[256];
  Chunk c = chunks[gid];
  df s2 = df_zero(), sadj = df_zero(), sneg = df_zero(), spos = df_zero(), s4 = df_zero();
  for (uint i = c.begin + lid; i < c.end; i += tsz) {
    float v = x[i]; float q = v * v;
    s2 = df_add_f(s2, q); s4 = df_add_f(s4, q * q);
    if (v < 0.0f) sneg = df_add_f(sneg, q); else if (v > 0.0f) spos = df_add_f(spos, q);
    if (i > c.begin) sadj = df_add_f(sadj, fabs(v) * fabs(x[i - 1]));
  }
  df r2 = tg_reduce(s2, sh, lid, tsz); threadgroup_barrier(mem_flags::mem_threadgroup);
  df radj = tg_reduce(sadj, sh, lid, tsz); threadgroup_barrier(mem_flags::mem_threadgroup);
  df rneg = tg_reduce(sneg, sh, lid, tsz); threadgroup_barrier(mem_flags::mem_threadgroup);
  df rpos = tg_reduce(spos, sh, lid, tsz); threadgroup_barrier(mem_flags::mem_threadgroup);
  df r4 = tg_reduce(s4, sh, lid, tsz);
  if (lid == 0) { out[gid * 5 + 0] = df_pack(r2); out[gid * 5 + 1] = df_pack(radj); out[gid * 5 + 2] = df_pack(rneg); out[gid * 5 + 3] = df_pack(rpos); out[gid * 5 + 4] = df_pack(r4); }
}

// Single-statistic variants (for the fusion ablation): op 0 = Σx², 1 = Σ|x||x_{-1}|, 2 = Σx² 1{x<0}, 3 = Σx² 1{x>0}, 4 = Σx⁴
kernel void segment_reduce(device const float* x [[buffer(0)]], device const Chunk* chunks [[buffer(1)]], device float2* out [[buffer(2)]], constant uint& op [[buffer(3)]],
                           uint gid [[threadgroup_position_in_grid]], uint lid [[thread_position_in_threadgroup]], uint tsz [[threads_per_threadgroup]]) {
  threadgroup df sh[256];
  Chunk c = chunks[gid]; df acc = df_zero();
  for (uint i = c.begin + lid; i < c.end; i += tsz) {
    float v = x[i]; float q = v * v; float t = 0.0f;
    switch (op) { case 0: t = q; break; case 1: t = (i > c.begin) ? fabs(v) * fabs(x[i - 1]) : 0.0f; break; case 2: t = v < 0.0f ? q : 0.0f; break; case 3: t = v > 0.0f ? q : 0.0f; break; default: t = q * q; }
    acc = df_add_f(acc, t);
  }
  df r = tg_reduce(acc, sh, lid, tsz);
  if (lid == 0) out[gid] = df_pack(r);
}

// ---- 2. rolling sum of squares (direct window, O(w) per output) ---------------------------------------------------
// out[i] = Σ_{j=i-w+1..i} x_j² within the segment; NaN where the window is incomplete. One thread per element.
kernel void rolling_sumsq(device const float* x [[buffer(0)]], device const uint* seg_begin [[buffer(1)]], device const uint* seg_of [[buffer(2)]], device float* out [[buffer(3)]],
                          constant uint& w [[buffer(4)]], constant uint& n [[buffer(5)]], uint i [[thread_position_in_grid]]) {
  if (i >= n) return;
  uint b = seg_begin[seg_of[i]];
  if (i + 1 < b + w) { out[i] = NAN; return; }
  df acc = df_zero();
  for (uint j = i + 1 - w; j <= i; ++j) { float v = x[j]; acc = df_add_f(acc, v * v); }
  out[i] = acc.hi + acc.lo;
}

// ---- 3. regression sufficient statistics (X'X upper triangle, X'y, y'y) per chunk ------------------------------
// X row-major (rows × k), k ≤ 8. One threadgroup per chunk; each thread accumulates its rows in double-float registers,
// then a deterministic tree reduction per statistic.
#define RAID_MAXK 8
kernel void regression_stats(device const float* X [[buffer(0)]], device const float* y [[buffer(1)]], device const Chunk* chunks [[buffer(2)]], device float2* out [[buffer(3)]],
                             constant uint& k [[buffer(4)]], uint gid [[threadgroup_position_in_grid]], uint lid [[thread_position_in_threadgroup]], uint tsz [[threads_per_threadgroup]]) {
  threadgroup df sh[256];
  Chunk c = chunks[gid];
  const uint ntri = k * (k + 1) / 2; const uint nstat = ntri + k + 1;
  df acc[RAID_MAXK * (RAID_MAXK + 1) / 2 + RAID_MAXK + 1];
  for (uint s = 0; s < nstat; ++s) acc[s] = df_zero();
  for (uint r = c.begin + lid; r < c.end; r += tsz) {
    float xr[RAID_MAXK]; for (uint a = 0; a < k; ++a) xr[a] = X[r * k + a];
    float yr = y[r]; uint s = 0;
    for (uint a = 0; a < k; ++a) for (uint b = a; b < k; ++b) { acc[s] = df_add_f(acc[s], xr[a] * xr[b]); ++s; }
    for (uint a = 0; a < k; ++a) { acc[s] = df_add_f(acc[s], xr[a] * yr); ++s; }
    acc[s] = df_add_f(acc[s], yr * yr);
  }
  for (uint s = 0; s < nstat; ++s) { df r = tg_reduce(acc[s], sh, lid, tsz); if (lid == 0) out[gid * nstat + s] = df_pack(r); threadgroup_barrier(mem_flags::mem_threadgroup); }
}

// ---- 4. batched GARCH-family negative log-likelihood ---------------------------------------------------------------
// One thread per (series, parameter vector). params layout per thread (stride 8): [mu, omega, alpha, gamma, beta, nu,
// kconst, eabs] — gamma=0 for GARCH, nu<=2 → Gaussian; kconst = density constant and eabs = E|z| are computed in FP64
// on the host (Metal has no lgamma). var0 = FP64-computed initial variance passed from the host (mixed precision).
// model: 0 garch, 1 gjr, 2 egarch. Returns NLL (double-float summed) as float; +INF for invalid parameters.
struct GarchJob { uint seg; uint pidx; uint pad0; uint pad1; };
kernel void garch_nll(device const float* r [[buffer(0)]], device const uint* seg_begin [[buffer(1)]], device const uint* seg_end [[buffer(2)]], device const float* params [[buffer(3)]],
                      device const float* var0 [[buffer(4)]], device const GarchJob* jobs [[buffer(5)]], device float* out [[buffer(6)]], constant uint& model [[buffer(7)]], constant uint& njobs [[buffer(8)]],
                      uint t [[thread_position_in_grid]]) {
  if (t >= njobs) return;
  GarchJob j = jobs[t]; const uint b = seg_begin[j.seg], e = seg_end[j.seg]; const uint n = e - b;
  const float mu = params[j.pidx * 8 + 0], omega = params[j.pidx * 8 + 1], alpha = params[j.pidx * 8 + 2], gamma = params[j.pidx * 8 + 3], beta = params[j.pidx * 8 + 4], nu = params[j.pidx * 8 + 5];
  const float kconst = params[j.pidx * 8 + 6], Eabs = params[j.pidx * 8 + 7];
  const bool student = nu > 2.0f;
  bool valid = true;
  if (model == 0u) valid = omega > 0.0f && alpha >= 0.0f && beta >= 0.0f && alpha + beta < 1.0f;
  else if (model == 1u) valid = omega > 0.0f && alpha >= 0.0f && beta >= 0.0f && alpha + gamma >= 0.0f && alpha + beta + 0.5f * gamma < 1.0f;
  else valid = fabs(beta) < 1.0f;
  if (!valid || n < 2) { out[t] = INFINITY; return; }
  float s2 = var0[j.seg];
  df nll = df_zero();
  float e_prev = r[b] - mu;
  for (uint i = 0; i < n; ++i) {
    if (i > 0) {
      if (model == 2u) { float z = e_prev / sqrt(s2); float ls = omega + alpha * (fabs(z) - Eabs) + gamma * z + beta * log(s2); s2 = exp(ls); }
      else { float e2 = e_prev * e_prev; s2 = omega + (alpha + (e_prev < 0.0f ? gamma : 0.0f)) * e2 + beta * s2; }
      if (!(s2 > 0.0f) || !isfinite(s2)) { out[t] = INFINITY; return; }
    }
    float ev = r[b + i] - mu; float ll;
    if (student) ll = kconst - 0.5f * log(s2) - 0.5f * (nu + 1.0f) * log(1.0f + ev * ev / ((nu - 2.0f) * s2));   // no log1p in MSL; argument is O(1)
    else ll = kconst - 0.5f * log(s2) - 0.5f * ev * ev / s2;
    nll = df_add_f(nll, -ll);
    e_prev = ev;
  }
  out[t] = nll.hi + nll.lo;
}