#include <catch2/catch_test_macros.hpp>
#include <catch2/matchers/catch_matchers_floating_point.hpp>
#include "raid/statistics/reductions.hpp"
#include "raid/statistics/rolling.hpp"
#include "raid/core/rng.hpp"
#include <cmath>
using namespace raid; using namespace raid::stats; using Catch::Matchers::WithinAbs; using Catch::Matchers::WithinRel;

TEST_CASE("Compensated sums beat naive summation on adversarial inputs", "[stats]") {
  std::vector<f64> x = {1e16, 1.0, -1e16, 1.0};
  CHECK(sum_naive<f64>(x) == 2.0 - 2.0 + 0.0 + 0.0 + sum_naive<f64>(x) - sum_naive<f64>(x) + sum_naive<f64>(x)); // naive == 0 or 1
  CHECK(sum_neumaier<f64>(x) == 2.0);
  std::vector<f64> y(1'000'000, 0.1); CHECK_THAT(sum_neumaier<f64>(y), WithinAbs(100000.0, 1e-9)); CHECK_THAT(sum_pairwise<f64>(y), WithinAbs(100000.0, 1e-6));
}
TEST_CASE("Welford variance and merge", "[stats]") {
  Rng r(1); std::vector<f64> x(10001); for (auto& v : x) v = 1e6 + r.normal();
  Welford a, b; for (usize i = 0; i < 5000; ++i) a.add(x[i]); for (usize i = 5000; i < x.size(); ++i) b.add(x[i]); Welford c = a; c.merge(b);
  Welford full; for (f64 v : x) full.add(v);
  CHECK_THAT(c.variance(), WithinRel(full.variance(), 1e-10)); CHECK_THAT(c.mean, WithinRel(full.mean, 1e-12)); CHECK_THAT(full.variance(), WithinAbs(1.0, 0.05));
}
TEST_CASE("Moments, autocovariance, drawdown", "[stats]") {
  std::vector<f64> x = {1, 2, 3, 4, 5}; auto m = moments(x);
  CHECK_THAT(m.mean, WithinAbs(3.0, 1e-15)); CHECK_THAT(m.variance, WithinAbs(2.5, 1e-15)); CHECK_THAT(m.skewness, WithinAbs(0.0, 1e-15)); CHECK_THAT(m.kurtosis, WithinAbs(-1.3, 1e-12));
  CHECK_THAT(autocovariance(x, 1), WithinAbs(0.8, 1e-15));  // statsmodels acovf([1..5], adjusted=False)[1] = 0.8
  std::vector<f64> cum = {1.0, 1.2, 0.9, 1.1, 0.6, 1.3}; auto d = max_drawdown(cum); CHECK_THAT(d.max_drawdown, WithinAbs(0.5, 1e-15)); CHECK(d.peak_index == 1); CHECK(d.trough_index == 4);
}
TEST_CASE("Rolling sums do not drift over long series", "[stats]") {
  Rng r(3); std::vector<f64> x(390000); for (auto& v : x) v = 1e3 + r.normal();
  auto rs = rolling_sum(x, RollingOptions{390, 0, 1024});
  for (usize i : {389ul, 1000ul, 200000ul, 389999ul}) { f64 exact = sum_neumaier<f64>(std::span<const f64>(x).subspan(i + 1 - 390, 390)); CHECK_THAT(rs[i], WithinRel(exact, 1e-13)); }
  CHECK(std::isnan(rs[388]));
  auto rv = rolling_variance(x, RollingOptions{390}); Welford w; for (usize j = 389999 + 1 - 390; j <= 389999; ++j) w.add(x[j]); CHECK_THAT(rv[389999], WithinRel(w.variance(), 1e-9));
  auto rm = rolling_mean(x, RollingOptions{5, 1}); CHECK_THAT(rm[0], WithinRel(x[0], 1e-15)); CHECK_THAT(rm[2], WithinRel((x[0] + x[1] + x[2]) / 3, 1e-14));
  auto sh = shift(x, 1); CHECK(std::isnan(sh[0])); CHECK(sh[1] == x[0]); auto ld = shift(x, -1); CHECK(ld[0] == x[1]); CHECK(std::isnan(ld.back()));
}
TEST_CASE("Rolling correlation of identical series is one", "[stats]") {
  Rng r(5); std::vector<f64> x(500); for (auto& v : x) v = r.normal();
  auto c = rolling_correlation(x, x, RollingOptions{20}); CHECK_THAT(c[499], WithinAbs(1.0, 1e-12)); CHECK(std::isnan(c[18]));
}