#include "raid/estimators/realized.hpp"
#include "raid/core/error.hpp"
#include "raid/statistics/reductions.hpp"
#include "raid/core/thread_pool.hpp"
#include <mutex>
#include <cmath>
#include <numbers>

namespace raid::est {

f64 mu_p(f64 p) { return std::pow(2.0, p / 2.0) * std::tgamma((p + 1.0) / 2.0) / std::tgamma(0.5); }

RealizedMeasures realized_measures_batch(const backend::Batch& b, const RealizedOptions& opt) {
  RealizedMeasures m; const usize S = b.n_series();
  std::string reason; auto be = backend::resolve(opt.plan, &reason);
  backend::KernelStats st1, st2, st3, st4; std::vector<f64> adj, s4; bool fused = false;
#ifdef RAID_ENABLE_METAL
  if (opt.fuse_kernels) fused = backend::realized_fused(*be, b, m.rv, adj, m.rsv_neg, m.rsv_pos, s4, &st1);
#endif
  if (!fused) {
    m.rv = be->sum_squares(b, &st1);
    adj = be->sum_abs_adjacent(b, &st2);
    be->semivariances(b, m.rsv_neg, m.rsv_pos, &st3);
    s4 = be->sum_fourth(b, &st4);
  }
  m.stats = st1; m.stats.kernel = fused ? "realized_fused" : "realized_measures(4 kernels)"; m.stats.wall_seconds = st1.wall_seconds + st2.wall_seconds + st3.wall_seconds + st4.wall_seconds;
  m.stats.device_seconds = st1.device_seconds + st2.device_seconds + st3.device_seconds + st4.device_seconds; m.stats.note = reason + (m.stats.note.empty() ? "" : " | " + m.stats.note);
  m.bpv.resize(S); m.rq.resize(S); m.jump.resize(S); m.tripv.resize(S); m.n_returns.resize(S); m.daily_return.resize(S);
  const f64 mu43 = mu_p(4.0 / 3.0);
  if (opt.tripower) {
    ThreadPool::global().parallel_for(S, [&](usize s0, usize s1) { std::vector<f64> p43; for (usize s = s0; s < s1; ++s) { const usize M = b.size(s); auto x = b.series(s);
        if (M <= 2) { m.tripv[s] = std::numeric_limits<f64>::quiet_NaN(); continue; }
        p43.resize(M); for (usize j = 0; j < M; ++j) { const f64 a = std::abs(x[j]); p43[j] = std::cbrt(a * a * a * a); }   // |x|^{4/3} = cbrt(x⁴)
        f64 acc = 0, c = 0; for (usize j = 2; j < M; ++j) { const f64 v = p43[j] * p43[j - 1] * p43[j - 2], t = acc + v; if (std::abs(acc) >= v) c += (acc - t) + v; else c += (v - t) + acc; acc = t; }
        const f64 fs3 = opt.bpv_finite_sample ? static_cast<f64>(M) / static_cast<f64>(M - 2) : 1.0; m.tripv[s] = fs3 * (acc + c) / (mu43 * mu43 * mu43); } });
  } else m.tripv.assign(S, std::numeric_limits<f64>::quiet_NaN());
  std::mutex warn_m;
  ThreadPool::global().parallel_for(S, [&](usize s0, usize s1) { for (usize s = s0; s < s1; ++s) {
    const usize M = b.size(s); m.n_returns[s] = M;
    auto x = b.series(s);
    m.daily_return[s] = stats::sum_neumaier(x);
    const f64 fs = (opt.bpv_finite_sample && M > 1) ? static_cast<f64>(M) / static_cast<f64>(M - 1) : 1.0;
    m.bpv[s] = M > 1 ? (1.0 / (kMu1 * kMu1)) * fs * adj[s] : std::numeric_limits<f64>::quiet_NaN();
    m.rq[s] = M > 0 ? (static_cast<f64>(M) / 3.0) * s4[s] : std::numeric_limits<f64>::quiet_NaN();
    m.jump[s] = std::max(m.rv[s] - m.bpv[s], 0.0);
    if (M < 2) { std::lock_guard<std::mutex> lk(warn_m); m.warnings.push_back("segment " + std::to_string(s) + ": fewer than 2 returns"); }
    if (opt.annualize) { m.rv[s] *= opt.periods_per_year; m.bpv[s] *= opt.periods_per_year; m.rsv_neg[s] *= opt.periods_per_year; m.rsv_pos[s] *= opt.periods_per_year; m.jump[s] *= opt.periods_per_year; m.tripv[s] *= opt.periods_per_year; m.rq[s] *= opt.periods_per_year * opt.periods_per_year; }
  } });
  return m;
}

RealizedMeasures realized_measures(const data::IntradayReturns& ir, const RealizedOptions& opt) {
  backend::Batch b{ir.r, ir.session_offsets};
  auto m = realized_measures_batch(b, opt);
  m.sessions = ir.session_dates;
  return m;
}

RealizedCov realized_covariance(const Eigen::Ref<const Eigen::MatrixXd>& R) {
  RAID_REQUIRE(R.rows() >= 1 && R.cols() >= 1, insufficient_samples, "realized_covariance: empty");
  RealizedCov out; out.n_obs = static_cast<usize>(R.rows());
  out.cov = R.transpose() * R;   // Σ_j r_j r_j' (no demeaning: realized covariance of returns)
  const Eigen::Index N = R.cols(); out.corr.resize(N, N);
  for (Eigen::Index a = 0; a < N; ++a) for (Eigen::Index c = 0; c < N; ++c) { const f64 d = std::sqrt(out.cov(a, a) * out.cov(c, c)); out.corr(a, c) = d > 0 ? out.cov(a, c) / d : std::numeric_limits<f64>::quiet_NaN(); }
  return out;
}

}  // namespace raid::est