#include "raid/statistics/bootstrap.hpp"
#include "raid/core/error.hpp"
#include "raid/core/thread_pool.hpp"
#include "raid/statistics/reductions.hpp"
#include <algorithm>
#include <chrono>
#include <cmath>

namespace raid::stats {

namespace {
f64 block_len(usize n, const BootstrapOptions& o) { return o.mean_block > 0 ? o.mean_block : std::max(1.0, std::round(std::cbrt(static_cast<f64>(n)))); }
f64 quantile(std::vector<f64> v, f64 q) { v.erase(std::remove_if(v.begin(), v.end(), [](f64 x) { return !std::isfinite(x); }), v.end()); if (v.empty()) return std::numeric_limits<f64>::quiet_NaN(); std::sort(v.begin(), v.end()); const f64 pos = q * static_cast<f64>(v.size() - 1); const usize i = static_cast<usize>(pos); const f64 fr = pos - static_cast<f64>(i); return i + 1 < v.size() ? v[i] * (1 - fr) + v[i + 1] * fr : v[i]; }
}

std::vector<u32> bootstrap_indices(usize n, const BootstrapOptions& o, usize replica) {
  RAID_REQUIRE(n >= 2, insufficient_samples, "bootstrap: n >= 2");
  Rng rng(o.seed, 1000003ULL + replica); std::vector<u32> idx(n); const f64 b = block_len(n, o);
  if (o.scheme == BootstrapScheme::iid) { for (usize i = 0; i < n; ++i) idx[i] = static_cast<u32>(rng.below(n)); return idx; }
  usize i = 0;
  if (o.scheme == BootstrapScheme::circular_block) { const usize bl = std::max<usize>(1, static_cast<usize>(std::llround(b))); while (i < n) { const usize start = rng.below(n); for (usize j = 0; j < bl && i < n; ++j) idx[i++] = static_cast<u32>((start + j) % n); } return idx; }
  const f64 p = 1.0 / b;  // stationary
  while (i < n) { const usize start = rng.below(n); const usize len = static_cast<usize>(rng.geometric(p)); for (usize j = 0; j < len && i < n; ++j) idx[i++] = static_cast<u32>((start + j) % n); }
  return idx;
}

BootstrapResult bootstrap(std::span<const f64> x, const std::function<f64(std::span<const f64>)>& stat, const BootstrapOptions& o) {
  const auto t0 = std::chrono::steady_clock::now(); BootstrapResult r; const usize n = x.size(); r.block_length = block_len(n, o); r.point = stat(x); r.statistics.assign(o.replicas, std::numeric_limits<f64>::quiet_NaN());
  ThreadPool pool(o.threads);
  pool.parallel_for(o.replicas, [&](usize b0, usize b1) { std::vector<f64> xs(n); for (usize b = b0; b < b1; ++b) { auto idx = bootstrap_indices(n, o, b); for (usize i = 0; i < n; ++i) xs[i] = x[idx[i]]; try { r.statistics[b] = stat(xs); } catch (...) { r.statistics[b] = std::numeric_limits<f64>::quiet_NaN(); } } });
  std::vector<f64> ok; for (f64 v : r.statistics) if (std::isfinite(v)) ok.push_back(v); else ++r.failures;
  if (ok.size() >= 2) { Welford w; for (f64 v : ok) w.add(v); r.se = std::sqrt(w.variance()); r.bias = w.mean - r.point; const f64 a = (1 - o.ci_level) / 2; r.ci_lo = quantile(ok, a); r.ci_hi = quantile(ok, 1 - a); }
  r.seconds = std::chrono::duration<f64>(std::chrono::steady_clock::now() - t0).count(); return r;
}

BootstrapVecResult bootstrap_rows(usize n, const std::function<std::vector<f64>(std::span<const u32>)>& stat, const BootstrapOptions& o) {
  const auto t0 = std::chrono::steady_clock::now(); BootstrapVecResult r; r.block_length = block_len(n, o);
  std::vector<u32> id(n); for (usize i = 0; i < n; ++i) id[i] = static_cast<u32>(i); r.point = stat(id); const usize k = r.point.size();
  r.statistics.assign(o.replicas, std::vector<f64>(k, std::numeric_limits<f64>::quiet_NaN())); ThreadPool pool(o.threads);
  pool.parallel_for(o.replicas, [&](usize b0, usize b1) { for (usize b = b0; b < b1; ++b) { auto idx = bootstrap_indices(n, o, b); try { auto s = stat(idx); if (s.size() == k) r.statistics[b] = s; } catch (...) {} } });
  r.se.assign(k, 0); r.ci_lo.assign(k, 0); r.ci_hi.assign(k, 0); const f64 a = (1 - o.ci_level) / 2;
  for (usize j = 0; j < k; ++j) { std::vector<f64> col; for (auto& s : r.statistics) if (std::isfinite(s[j])) col.push_back(s[j]); if (j == 0) r.failures = o.replicas - col.size(); if (col.size() >= 2) { Welford w; for (f64 v : col) w.add(v); r.se[j] = std::sqrt(w.variance()); r.ci_lo[j] = quantile(col, a); r.ci_hi[j] = quantile(col, 1 - a); } }
  r.seconds = std::chrono::duration<f64>(std::chrono::steady_clock::now() - t0).count(); return r;
}

BootstrapResult bootstrap_rv(std::span<const f64> r, const BootstrapOptions& o) { return bootstrap(r, [](std::span<const f64> x) { return sum_squares(x); }, o); }

}  // namespace raid::stats