#include "raid/statistics/rolling.hpp"
#include "raid/statistics/reductions.hpp"
#include "raid/core/error.hpp"
#include <cmath>

namespace raid::stats {

namespace {
struct RunningSum {
  f64 s{0}, c{0};
  void add(f64 v) { const f64 t = s + v; if (std::abs(s) >= std::abs(v)) c += (s - t) + v; else c += (v - t) + s; s = t; }
  f64 value() const { return s + c; }
  void reset(f64 exact) { s = exact; c = 0; }
};
inline usize minp(const RollingOptions& o) { return o.min_periods == 0 ? o.window : o.min_periods; }
}  // namespace

std::vector<f64> rolling_sum(std::span<const f64> x, const RollingOptions& o) {
  RAID_REQUIRE(o.window >= 1, invalid_argument, "window must be >= 1");
  const usize n = x.size(), w = o.window, mp = minp(o);
  std::vector<f64> out(n, kNaN);
  RunningSum rs; usize since_resum = 0;
  for (usize i = 0; i < n; ++i) {
    rs.add(x[i]);
    if (i >= w) rs.add(-x[i - w]);
    if (++since_resum >= o.resum_every) { const usize b = i + 1 >= w ? i + 1 - w : 0; rs.reset(sum_neumaier(x.subspan(b, i + 1 - b))); since_resum = 0; }
    const usize count = i + 1 < w ? i + 1 : w;
    if (count >= mp) out[i] = rs.value();
  }
  return out;
}

std::vector<f64> rolling_mean(std::span<const f64> x, const RollingOptions& o) {
  auto s = rolling_sum(x, o);
  for (usize i = 0; i < s.size(); ++i) { const usize count = i + 1 < o.window ? i + 1 : o.window; s[i] /= static_cast<f64>(count); }
  return s;
}

std::vector<f64> rolling_variance(std::span<const f64> x, const RollingOptions& o, int ddof) {
  // Two accumulators (sum, sum of squares about a per-window shift) would still cancel for large means;
  // we recompute each window by Welford from scratch when the window is small (<= 64), otherwise use the
  // shifted-sums formula with periodic exact re-summation. Both are FP64 references; speed is not the goal here.
  const usize n = x.size(), w = o.window, mp = minp(o);
  std::vector<f64> out(n, kNaN);
  if (w <= 64) {
    for (usize i = 0; i < n; ++i) {
      const usize b = i + 1 >= w ? i + 1 - w : 0; const usize count = i + 1 - b;
      if (count < mp || count <= static_cast<usize>(ddof)) continue;
      Welford wf; for (usize j = b; j <= i; ++j) wf.add(x[j]);
      out[i] = wf.variance(ddof);
    }
    return out;
  }
  const f64 shift_c = x.empty() ? 0.0 : x[0];
  RunningSum s1, s2; usize since = 0;
  for (usize i = 0; i < n; ++i) {
    const f64 d = x[i] - shift_c; s1.add(d); s2.add(d * d);
    if (i >= w) { const f64 e = x[i - w] - shift_c; s1.add(-e); s2.add(-e * e); }
    if (++since >= o.resum_every) {
      const usize b = i + 1 >= w ? i + 1 - w : 0; f64 a = 0, q = 0;
      std::vector<f64> t1(i + 1 - b), t2(i + 1 - b);
      for (usize j = b; j <= i; ++j) { t1[j - b] = x[j] - shift_c; t2[j - b] = t1[j - b] * t1[j - b]; }
      a = sum_neumaier(std::span<const f64>(t1)); q = sum_neumaier(std::span<const f64>(t2)); s1.reset(a); s2.reset(q); since = 0;
    }
    const usize count = i + 1 < w ? i + 1 : w;
    if (count >= mp && count > static_cast<usize>(ddof)) {
      const f64 cnt = static_cast<f64>(count), m = s1.value() / cnt;
      out[i] = std::max(0.0, (s2.value() - cnt * m * m) / (cnt - static_cast<f64>(ddof)));
    }
  }
  return out;
}

std::vector<f64> rolling_covariance(std::span<const f64> x, std::span<const f64> y, const RollingOptions& o, int ddof) {
  RAID_REQUIRE(x.size() == y.size(), invalid_argument, "rolling_covariance: size mismatch");
  const usize n = x.size(), w = o.window, mp = minp(o);
  std::vector<f64> out(n, kNaN);
  for (usize i = 0; i < n; ++i) {
    const usize b = i + 1 >= w ? i + 1 - w : 0; const usize count = i + 1 - b;
    if (count < mp || count <= static_cast<usize>(ddof)) continue;
    f64 mx = sum_neumaier(x.subspan(b, count)) / static_cast<f64>(count), my = sum_neumaier(y.subspan(b, count)) / static_cast<f64>(count);
    f64 s = 0, c = 0;
    for (usize j = b; j <= i; ++j) { const f64 v = (x[j] - mx) * (y[j] - my), t = s + v; if (std::abs(s) >= std::abs(v)) c += (s - t) + v; else c += (v - t) + s; s = t; }
    out[i] = (s + c) / static_cast<f64>(count - static_cast<usize>(ddof));
  }
  return out;
}

std::vector<f64> rolling_correlation(std::span<const f64> x, std::span<const f64> y, const RollingOptions& o) {
  auto cxy = rolling_covariance(x, y, o, 1); auto vx = rolling_variance(x, o, 1); auto vy = rolling_variance(y, o, 1);
  for (usize i = 0; i < cxy.size(); ++i) cxy[i] = (vx[i] > 0 && vy[i] > 0) ? cxy[i] / std::sqrt(vx[i] * vy[i]) : kNaN;
  return cxy;
}

std::vector<f64> expanding_mean(std::span<const f64> x, usize min_periods) {
  std::vector<f64> out(x.size(), kNaN); Welford w;
  for (usize i = 0; i < x.size(); ++i) { w.add(x[i]); if (w.n >= min_periods) out[i] = w.mean; }
  return out;
}
std::vector<f64> expanding_variance(std::span<const f64> x, usize min_periods, int ddof) {
  std::vector<f64> out(x.size(), kNaN); Welford w;
  for (usize i = 0; i < x.size(); ++i) { w.add(x[i]); if (w.n >= min_periods) out[i] = w.variance(ddof); }
  return out;
}
std::vector<f64> shift(std::span<const f64> x, long k) {
  std::vector<f64> out(x.size(), kNaN);
  for (usize i = 0; i < x.size(); ++i) { const long j = static_cast<long>(i) - k; if (j >= 0 && j < static_cast<long>(x.size())) out[i] = x[static_cast<usize>(j)]; }
  return out;
}

}  // namespace raid::stats