#include "raid/estimators/har.hpp"
#include "raid/core/error.hpp"
#include "raid/statistics/covariance.hpp"
#include <chrono>
#include <cmath>
#include <algorithm>

namespace raid::est {

HarDesign har_design(std::span<const f64> rv, const HarSpec& spec, std::span<const f64> rq, const Eigen::MatrixXd& exog, const std::vector<std::string>& exog_names) {
  const usize T = rv.size(); RAID_REQUIRE(!spec.lags.empty() && spec.horizon >= 1, invalid_argument, "har: lags and horizon");
  const usize maxlag = *std::max_element(spec.lags.begin(), spec.lags.end()); const usize h = static_cast<usize>(spec.horizon);
  RAID_REQUIRE(T > maxlag + h + 2, insufficient_samples, "har: need T > maxlag + h + 2");
  for (f64 v : rv) RAID_REQUIRE(std::isfinite(v) && (!spec.log || v > 0), data_quality, "har: RV must be finite (and > 0 for log-HAR)");
  if (spec.harq) RAID_REQUIRE(rq.size() == T, invalid_argument, "harq: RQ series required with the same length");
  const usize nex = exog.size() ? static_cast<usize>(exog.cols()) : 0; if (nex) RAID_REQUIRE(static_cast<usize>(exog.rows()) == T, invalid_argument, "har-x: exog rows must equal T");
  std::vector<f64> z(T); for (usize t = 0; t < T; ++t) z[t] = spec.log ? std::log(rv[t]) : rv[t];
  // prefix sums for the window means (FP64, Neumaier)
  std::vector<f64> cs(T + 1, 0.0); { f64 s = 0, c = 0; for (usize t = 0; t < T; ++t) { const f64 v = z[t], tt = s + v; if (std::abs(s) >= std::abs(v)) c += (s - tt) + v; else c += (v - tt) + s; s = tt; cs[t + 1] = s + c; } }
  auto wmean = [&](usize t, usize k) { f64 s = 0; for (usize j = 0; j < k; ++j) s += z[t - j]; return s / static_cast<f64>(k); };  // direct sum (exact enough, k ≤ 22)
  const usize first = maxlag - 1, last = T - 1 - h;  // origins t with full lags and a full target
  HarDesign d; d.first_t = first; d.n = last - first + 1; const usize k = 1 + spec.lags.size() + (spec.harq ? 1 : 0) + nex;
  d.X.resize(static_cast<Eigen::Index>(d.n), static_cast<Eigen::Index>(k)); d.y.resize(static_cast<Eigen::Index>(d.n));
  d.names.push_back("const"); for (usize l : spec.lags) d.names.push_back((spec.log ? "logrv_" : "rv_") + std::to_string(l)); if (spec.harq) d.names.push_back("sqrt_rq_x_rv"); for (usize j = 0; j < nex; ++j) d.names.push_back(j < exog_names.size() ? exog_names[j] : "x" + std::to_string(j));
  for (usize t = first; t <= last; ++t) {
    const Eigen::Index i = static_cast<Eigen::Index>(t - first); Eigen::Index c = 0; d.X(i, c++) = 1.0;
    for (usize l : spec.lags) d.X(i, c++) = wmean(t, l);
    if (spec.harq) d.X(i, c++) = std::sqrt(std::max(rq[t], 0.0)) * rv[t];
    for (usize j = 0; j < nex; ++j) d.X(i, c++) = exog(static_cast<Eigen::Index>(t), static_cast<Eigen::Index>(j));
    if (spec.target == HarTarget::single) d.y[i] = z[t + h]; else { f64 s = 0; for (usize j = 1; j <= h; ++j) s += spec.log ? rv[t + j] : z[t + j]; s /= static_cast<f64>(h); d.y[i] = spec.log ? std::log(s) : s; }
  }
  return d;
}

Estimate har_fit(std::span<const f64> rv, const HarSpec& spec, std::span<const f64> rq, const Eigen::MatrixXd& exog, const std::vector<std::string>& exog_names) {
  auto d = har_design(rv, spec, rq, exog, exog_names);
  OlsOptions o; o.cov = spec.cov; o.names = d.names;
  if (o.cov.type == stats::CovType::hac_bartlett && o.cov.hac_lags < 0) o.cov.hac_lags = std::max(spec.horizon - 1, stats::default_hac_lags(d.n));
  auto r = ols(d.X, d.y, o); Estimate e = r.estimate; e.model = spec.model_key();
  e.extra["horizon"] = spec.horizon; e.extra["target"] = spec.target == HarTarget::average ? "average" : "single"; e.extra["lags"] = spec.lags; e.extra["first_origin"] = d.first_t; e.extra["log"] = spec.log; e.extra["harq"] = spec.harq;
  if (spec.log) { f64 s = 0; for (Eigen::Index i = 0; i < r.residuals.size(); ++i) s += std::exp(r.residuals[i]); e.extra["smearing"] = s / static_cast<f64>(r.residuals.size()); }
  return e;
}

f64 har_forecast(const Estimate& fit, std::span<const f64> rv, const HarSpec& spec, std::span<const f64> rq, const Eigen::VectorXd& exog_last) {
  const usize T = rv.size(); const usize maxlag = *std::max_element(spec.lags.begin(), spec.lags.end()); RAID_REQUIRE(T >= maxlag, insufficient_samples, "har_forecast: not enough history");
  const usize t = T - 1; f64 yhat = fit.coefficients[0]; usize c = 1;
  for (usize l : spec.lags) { f64 s = 0; for (usize j = 0; j < l; ++j) s += spec.log ? std::log(rv[t - j]) : rv[t - j]; yhat += fit.coefficients[c++] * s / static_cast<f64>(l); }
  if (spec.harq) { RAID_REQUIRE(rq.size() == T, invalid_argument, "harq forecast needs RQ"); yhat += fit.coefficients[c++] * std::sqrt(std::max(rq[t], 0.0)) * rv[t]; }
  for (Eigen::Index j = 0; j < exog_last.size(); ++j) yhat += fit.coefficients[c++] * exog_last[j];
  return spec.log ? std::exp(yhat) : yhat;
}

HarRolling har_rolling(std::span<const f64> rv, const HarSpec& spec, usize window, usize step, std::span<const f64> rq) {
  const auto t0 = std::chrono::steady_clock::now(); HarRolling out; const usize T = rv.size(); const usize h = static_cast<usize>(spec.horizon);
  const usize maxlag = *std::max_element(spec.lags.begin(), spec.lags.end()); const usize minT = maxlag + h + 30;
  RAID_REQUIRE(window == 0 || window >= minT, insufficient_samples, "har_rolling: window too small"); RAID_REQUIRE(T > minT + h, insufficient_samples, "har_rolling: series too short");
  const usize k = 1 + spec.lags.size() + (spec.harq ? 1 : 0); std::vector<Eigen::VectorXd> coefs;
  for (usize end = (window ? window : minT); end + h <= T; end += step) {       // model uses rv[0..end) (or the last `window`), forecasts target at origin end−1
    const usize b = window ? end - window : 0; std::span<const f64> hist = rv.subspan(b, end - b); std::span<const f64> rqh = spec.harq ? rq.subspan(b, end - b) : std::span<const f64>();
    Estimate e; try { HarSpec s2 = spec; s2.cov.type = stats::CovType::classical; e = har_fit(hist, s2, rqh); } catch (const RaidError&) { continue; }
    const f64 fc = har_forecast(e, hist, spec, rqh); const usize t = end - 1;
    f64 actual; if (spec.target == HarTarget::single) actual = rv[t + h]; else { f64 s = 0; for (usize j = 1; j <= h; ++j) s += rv[t + j]; actual = s / static_cast<f64>(h); }
    out.origin.push_back(t); out.forecast.push_back(fc); out.actual.push_back(actual); coefs.push_back(Eigen::Map<const Eigen::VectorXd>(e.coefficients.data(), static_cast<Eigen::Index>(k))); ++out.n_fits;
  }
  out.coefficients.resize(static_cast<Eigen::Index>(k), static_cast<Eigen::Index>(coefs.size())); for (usize j = 0; j < coefs.size(); ++j) out.coefficients.col(static_cast<Eigen::Index>(j)) = coefs[j];
  out.seconds = std::chrono::duration<f64>(std::chrono::steady_clock::now() - t0).count(); return out;
}

HarBatchResult har_batch(std::span<const f64> rv_concat, std::span<const i64> offsets, const HarSpec& spec, const ExecutionPlan& plan) {
  HarBatchResult out; const usize S = offsets.size() - 1; std::string why; auto be = backend::resolve(plan, &why);
  // Build all designs (CPU, FP64), concatenate row-major for the backend, accumulate X'X / X'y / y'y per series, solve in FP64.
  std::vector<HarDesign> designs(S); usize k = 0, rows = 0; std::vector<i64> roff(S + 1, 0); std::vector<usize> ok(S, 1);
  for (usize s = 0; s < S; ++s) { try { designs[s] = har_design(rv_concat.subspan(static_cast<usize>(offsets[s]), static_cast<usize>(offsets[s + 1] - offsets[s])), spec); k = static_cast<usize>(designs[s].X.cols()); rows += designs[s].n; } catch (const RaidError&) { ok[s] = 0; } roff[s + 1] = static_cast<i64>(rows); }
  RAID_REQUIRE(k > 0, insufficient_samples, "har_batch: no series with enough data");
  std::vector<f64> X(rows * k), y(rows); for (usize s = 0; s < S; ++s) { if (!ok[s]) continue; const usize b = static_cast<usize>(roff[s]); for (usize i = 0; i < designs[s].n; ++i) { y[b + i] = designs[s].y[static_cast<Eigen::Index>(i)]; for (usize j = 0; j < k; ++j) X[(b + i) * k + j] = designs[s].X(static_cast<Eigen::Index>(i), static_cast<Eigen::Index>(j)); } }
  std::vector<f64> xtx, xty, yty; be->regression_stats(X, y, roff, k, xtx, xty, yty, &out.stats);
  out.fits.resize(S);
  for (usize s = 0; s < S; ++s) {
    if (!ok[s]) { out.fits[s].model = spec.model_key(); out.fits[s].convergence = Convergence{false, 0, 0, 0, 1, "insufficient data"}; continue; }
    SufficientStats ss; ss.xtx = Eigen::Map<const Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(xtx.data() + s * k * k, static_cast<Eigen::Index>(k), static_cast<Eigen::Index>(k)); ss.xty = Eigen::Map<const Eigen::VectorXd>(xty.data() + s * k, static_cast<Eigen::Index>(k)); ss.yty = yty[s]; ss.n = designs[s].n;
    OlsOptions o; o.names = designs[s].names; auto r = ols_from_stats(ss, o); out.fits[s] = r.estimate; out.fits[s].model = spec.model_key(); out.fits[s].backend = be->kind(); out.fits[s].precision = out.stats.precision;
    out.fits[s].extra["horizon"] = spec.horizon; out.fits[s].extra["note"] = "classical SE from sufficient statistics; use har_fit for HAC";
  }
  return out;
}

ForecastLoss forecast_loss(std::span<const f64> a, std::span<const f64> f) {
  RAID_REQUIRE(a.size() == f.size(), invalid_argument, "forecast_loss: size mismatch"); ForecastLoss l{0, 0, 0};
  for (usize i = 0; i < a.size(); ++i) { if (!(a[i] > 0) || !(f[i] > 0) || !std::isfinite(a[i]) || !std::isfinite(f[i])) continue; const f64 d = a[i] - f[i]; l.mse += d * d; const f64 q = a[i] / f[i]; l.qlike += q - std::log(q) - 1.0; ++l.n; }
  if (l.n) { l.mse /= static_cast<f64>(l.n); l.qlike /= static_cast<f64>(l.n); } return l;
}

}  // namespace raid::est