// Metal backend: accuracy vs the FP64 CPU reference, determinism, chunking, fallback. Skipped when Metal is absent.
#include <catch2/catch_test_macros.hpp>
#include <catch2/matchers/catch_matchers_floating_point.hpp>
#include "raid/backend/backend.hpp"
#include "raid/core/rng.hpp"
#include "raid/estimators/realized.hpp"
#include "raid/estimators/ols.hpp"
#include "raid/models/garch.hpp"
#include <cmath>
using namespace raid; using Catch::Matchers::WithinRel; using Catch::Matchers::WithinAbs;

#ifdef RAID_ENABLE_METAL
static std::shared_ptr<backend::Backend> metal() { auto m = backend::make_metal_backend(); return (m && m->available()) ? m : nullptr; }

static void make_batch(usize S, usize n, std::vector<f64>& v, std::vector<i64>& off, u64 seed = 1) {
  Rng r(seed); v.resize(S * n); off.resize(S + 1); for (usize s = 0; s <= S; ++s) off[s] = static_cast<i64>(s * n);
  for (auto& x : v) x = 1e-3 * r.student_t(5.0);
}

TEST_CASE("Metal realized accumulators match FP64 CPU within 1e-6 (fused and single kernels, chunked)", "[metal]") {
  auto m = metal(); if (!m) { SKIP("no Metal device"); }
  auto cpu = backend::make_cpu_backend();
  for (usize n : {390ul, 3900ul, 39000ul}) {
    std::vector<f64> v; std::vector<i64> off; make_batch(7, n, v, off, n); backend::Batch b{v, off};
    backend::KernelStats st; std::vector<f64> s2, adj, neg, pos, s4; REQUIRE(backend::realized_fused(*m, b, s2, adj, neg, pos, s4, &st));
    auto c2 = cpu->sum_squares(b), cadj = cpu->sum_abs_adjacent(b), c4 = cpu->sum_fourth(b); std::vector<f64> cn, cp; cpu->semivariances(b, cn, cp);
    for (usize s = 0; s < 7; ++s) { INFO("n=" << n << " s=" << s); CHECK_THAT(s2[s], WithinRel(c2[s], 1e-6)); CHECK_THAT(adj[s], WithinRel(cadj[s], 1e-6)); CHECK_THAT(neg[s], WithinRel(cn[s], 1e-6)); CHECK_THAT(pos[s], WithinRel(cp[s], 1e-6)); CHECK_THAT(s4[s], WithinRel(c4[s], 1e-5)); }
    CHECK(st.backend == Backend::metal); CHECK(st.precision == Precision::fp32); CHECK(st.device_seconds > 0);
    auto g2 = m->sum_squares(b); auto gadj = m->sum_abs_adjacent(b); for (usize s = 0; s < 7; ++s) { CHECK_THAT(g2[s], WithinRel(c2[s], 1e-6)); CHECK_THAT(gadj[s], WithinRel(cadj[s], 1e-6)); }
  }
}
TEST_CASE("Metal results are bitwise deterministic across runs", "[metal]") {
  auto m = metal(); if (!m) { SKIP("no Metal device"); }
  std::vector<f64> v; std::vector<i64> off; make_batch(50, 3900, v, off, 9); backend::Batch b{v, off};
  std::vector<f64> a2, aa, an, ap, a4, b2, ba, bn, bp, b4; backend::realized_fused(*m, b, a2, aa, an, ap, a4); backend::realized_fused(*m, b, b2, ba, bn, bp, b4);
  CHECK(a2 == b2); CHECK(aa == ba); CHECK(an == bn); CHECK(a4 == b4);
}
TEST_CASE("Metal rolling sum of squares and regression statistics match CPU", "[metal]") {
  auto m = metal(); if (!m) { SKIP("no Metal device"); } auto cpu = backend::make_cpu_backend();
  std::vector<f64> v; std::vector<i64> off; make_batch(3, 2000, v, off, 4); backend::Batch b{v, off};
  auto g = m->rolling_sum_squares(b, 22), c = cpu->rolling_sum_squares(b, 22);
  for (usize i = 0; i < v.size(); ++i) { if (std::isnan(c[i])) { CHECK(std::isnan(g[i])); continue; } CHECK_THAT(g[i], WithinRel(c[i], 1e-5)); }
  // regression: 5 series × 1000 rows × k=4
  Rng r(3); const usize S = 5, n = 1000, k = 4; std::vector<f64> X(S * n * k), y(S * n); std::vector<i64> ro(S + 1); for (usize s = 0; s <= S; ++s) ro[s] = static_cast<i64>(s * n);
  for (usize i = 0; i < S * n; ++i) { X[i * k] = 1.0; for (usize j = 1; j < k; ++j) X[i * k + j] = r.normal(); y[i] = 0.5 + X[i * k + 1] - 0.3 * X[i * k + 2] + r.normal(); }
  std::vector<f64> gx, gy, gyy, cx, cy, cyy; backend::KernelStats st; m->regression_stats(X, y, ro, k, gx, gy, gyy, &st); cpu->regression_stats(X, y, ro, k, cx, cy, cyy);
  for (usize i = 0; i < gx.size(); ++i) CHECK_THAT(gx[i], WithinAbs(cx[i], 1e-6 * std::max(1.0, std::abs(cx[i])) + 1e-9)); for (usize i = 0; i < gy.size(); ++i) CHECK_THAT(gy[i], WithinAbs(cy[i], 1e-6 * std::max(1.0, std::abs(cy[i])) + 1e-9)); for (usize s = 0; s < S; ++s) CHECK_THAT(gyy[s], WithinRel(cyy[s], 1e-6));
  // solve from GPU sufficient statistics and compare coefficients with the CPU QR solution
  Eigen::MatrixXd Xm(n, k); Eigen::VectorXd ym(n); for (usize i = 0; i < n; ++i) { ym[i] = y[i]; for (usize j = 0; j < k; ++j) Xm(i, j) = X[i * k + j]; }
  est::SufficientStats ss; ss.xtx = Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(gx.data(), k, k); ss.xty = Eigen::Map<Eigen::VectorXd>(gy.data(), k); ss.yty = gyy[0]; ss.n = n;
  auto from_gpu = est::ols_from_stats(ss); auto ref = est::ols(Xm, ym); for (usize j = 0; j < k; ++j) CHECK_THAT(from_gpu.estimate.coefficients[j], WithinAbs(ref.estimate.coefficients[j], 1e-5));
}
TEST_CASE("Metal batched GARCH likelihood matches FP64 within tolerance", "[metal]") {
  auto m = metal(); if (!m) { SKIP("no Metal device"); } auto cpu = backend::make_cpu_backend();
  Rng rng(11); const usize S = 4, n = 2000; std::vector<f64> v(S * n); std::vector<i64> off(S + 1); for (usize s = 0; s <= S; ++s) off[s] = static_cast<i64>(s * n);
  for (usize s = 0; s < S; ++s) { f64 s2 = 0.04 / 0.05, e = 0; for (usize t = 0; t < n; ++t) { s2 = 0.04 + 0.05 * e * e + 0.9 * s2; e = std::sqrt(s2) * rng.normal(); v[s * n + t] = 0.01 + e; } }
  backend::Batch b{v, off};
  const usize P = 16; std::vector<f64> params(S * P * 4); for (usize s = 0; s < S; ++s) for (usize q = 0; q < P; ++q) { f64* th = params.data() + (s * P + q) * 4; th[0] = 0.01; th[1] = 0.02 + 0.004 * q; th[2] = 0.03 + 0.005 * q; th[3] = 0.85 + 0.005 * q; }
  backend::KernelStats st; auto g = m->garch_nll_batch(b, params, P, 4, 0, 0.0, &st); auto c = cpu->garch_nll_batch(b, params, P, 4, 0, 0.0);
  for (usize j = 0; j < S * P; ++j) { INFO("job " << j); CHECK_THAT(g[j], WithinRel(c[j], 2e-5)); }
  // argmin over the grid must agree
  for (usize s = 0; s < S; ++s) { usize ag = 0, ac = 0; for (usize q = 1; q < P; ++q) { if (g[s * P + q] < g[s * P + ag]) ag = q; if (c[s * P + q] < c[s * P + ac]) ac = q; } CHECK(ag == ac); }
  // GJR with gamma and Student-t fixed nu
  std::vector<f64> p5(S * P * 5); for (usize s = 0; s < S; ++s) for (usize q = 0; q < P; ++q) { f64* th = p5.data() + (s * P + q) * 5; th[0] = 0.01; th[1] = 0.03; th[2] = 0.03; th[3] = 0.02 + 0.002 * q; th[4] = 0.88; }
  auto g5 = m->garch_nll_batch(b, p5, P, 5, 1, 7.0, &st); auto c5 = cpu->garch_nll_batch(b, p5, P, 5, 1, 7.0); for (usize j = 0; j < S * P; ++j) CHECK_THAT(g5[j], WithinRel(c5[j], 5e-5));
}
TEST_CASE("Metal invalid parameters give +inf like the CPU", "[metal]") {
  auto m = metal(); if (!m) { SKIP("no Metal device"); }
  std::vector<f64> v(300, 0.01); std::vector<i64> off = {0, 300}; backend::Batch b{v, off}; std::vector<f64> bad = {0.0, 0.1, 0.6, 0.5};
  auto g = m->garch_nll_batch(b, bad, 1, 4, 0, 0.0); CHECK(std::isinf(g[0]));
}
#endif

TEST_CASE("Backend resolution and fallback", "[backend]") {
  ExecutionPlan p; p.backend = Backend::cpu; std::string why; auto be = backend::resolve(p, &why); CHECK(be->kind() == Backend::cpu);
  p.backend = Backend::metal; p.allow_fallback = true; auto be2 = backend::resolve(p, &why); CHECK((be2->kind() == Backend::metal || be2->kind() == Backend::cpu));
  auto devs = backend::list_devices(); CHECK(!devs.empty()); CHECK(devs[0].kind == Backend::cpu); CHECK(devs[0].supports_fp64);
  for (auto& d : devs) if (d.kind == Backend::metal) CHECK_FALSE(d.supports_fp64);
}