// GARCH-family QML estimation on the CPU: BFGS on the unconstrained parameters, central-difference gradients,
// Armijo backtracking, and three covariance estimators (Hessian, OPG, QMLE sandwich = H^{-1} OPG H^{-1}).
#include "raid/models/garch.hpp"
#include "raid/core/error.hpp"
#include "raid/core/rng.hpp"
#include "raid/estimators/ols.hpp"
#include "raid/statistics/reductions.hpp"
#include <Eigen/Dense>
#include <chrono>
#include <cmath>
#include <numbers>
#include <cstdlib>
#include <cstdio>

namespace raid::models {

namespace {
using Vec = Eigen::VectorXd; using Mat = Eigen::MatrixXd;

// Per-observation log-likelihood contributions ℓ_t (returns +inf nll on invalid params).
bool per_obs_ll(std::span<const f64> r, std::span<const f64> th, GarchModel m, GarchDist d, std::vector<f64>& ll) {
  if (!garch_params_valid(th, m, d)) return false;
  std::vector<f64> s2; if (!garch_variance_path(r, th, m, d, s2)) return false;
  const f64 mu = th[0]; const usize n = r.size(); ll.resize(n);
  if (d == GarchDist::gaussian) { const f64 k = 0.5 * std::log(2.0 * std::numbers::pi); for (usize t = 0; t < n; ++t) { const f64 e = r[t] - mu; ll[t] = -(k + 0.5 * std::log(s2[t]) + 0.5 * e * e / s2[t]); } }
  else { const f64 nu = th.back(); const f64 k = std::lgamma((nu + 1) / 2) - std::lgamma(nu / 2) - 0.5 * std::log(std::numbers::pi * (nu - 2)); for (usize t = 0; t < n; ++t) { const f64 e = r[t] - mu; ll[t] = k - 0.5 * std::log(s2[t]) - 0.5 * (nu + 1) * std::log1p(e * e / ((nu - 2) * s2[t])); } }
  return true;
}

template <class F> Vec num_grad(const F& f, const Vec& x, f64 fx, f64 h0 = 1e-5) {
  Vec g(x.size());
  for (Eigen::Index i = 0; i < x.size(); ++i) { const f64 h = h0 * std::max(1.0, std::abs(x[i])); Vec xp = x, xm = x; xp[i] += h; xm[i] -= h; const f64 fp = f(xp), fm = f(xm);
    if (std::isfinite(fp) && std::isfinite(fm)) g[i] = (fp - fm) / (2 * h); else if (std::isfinite(fp)) g[i] = (fp - fx) / h; else if (std::isfinite(fm)) g[i] = (fx - fm) / h; else g[i] = 0; }
  return g;
}
template <class F> Mat num_hessian(const F& f, const Vec& x, f64 h0 = 1e-4) {
  const Eigen::Index p = x.size(); Mat H(p, p); const f64 fx = f(x);
  for (Eigen::Index i = 0; i < p; ++i) for (Eigen::Index j = i; j < p; ++j) {
    const f64 hi = h0 * std::max(1.0, std::abs(x[i])), hj = h0 * std::max(1.0, std::abs(x[j]));
    Vec a = x, b = x, c = x, dd = x; a[i] += hi; a[j] += hj; b[i] += hi; b[j] -= hj; c[i] -= hi; c[j] += hj; dd[i] -= hi; dd[j] -= hj;
    f64 v = (f(a) - f(b) - f(c) + f(dd)) / (4 * hi * hj);
    if (i == j) { Vec xp = x, xm = x; xp[i] += hi; xm[i] -= hi; v = (f(xp) - 2 * fx + f(xm)) / (hi * hi); }
    H(i, j) = v; H(j, i) = v; }
  return H;
}
}  // namespace

static Estimate garch_fit_once(std::span<const f64> r0, const GarchFitOptions& opt, int start_variant);
Estimate garch_fit(std::span<const f64> r0, const GarchFitOptions& opt) {
  // Up to 3 starting-value variants (documented in the registry): 0 = default, 1 = low persistence (0.80), 2 = high persistence (0.985) / opposite EGARCH γ sign.
  Estimate best; bool have = false; int tried = 0;
  for (int v = 0; v < 3; ++v) { Estimate e = garch_fit_once(r0, opt, v); ++tried; if (!have || (e.ok() && (!best.ok() || e.loss < best.loss))) { best = e; have = true; } if (best.ok() && opt.retry_starts == false) break; if (best.ok()) break; }
  best.extra["start_variants_tried"] = tried; return best;
}
static Estimate garch_fit_once(std::span<const f64> r0, const GarchFitOptions& opt, int start_variant) {
  const auto t0 = std::chrono::steady_clock::now();
  Estimate e; e.model = "garch." + to_string(opt.model) + "." + to_string(opt.dist); e.backend = Backend::cpu; e.precision = Precision::fp64;
  e.names = garch_param_names(opt.model, opt.dist); e.n_obs = r0.size();
  RAID_REQUIRE(r0.size() >= 50, insufficient_samples, "garch_fit: need at least 50 observations");
  for (f64 v : r0) RAID_REQUIRE(std::isfinite(v), data_quality, "garch_fit: non-finite return");
  // scaling (arch rescales to unit variance ×100-ish; we scale so that sample std ≈ 1 and unscale afterwards)
  const f64 sd = std::sqrt(stats::variance(r0, 0)); RAID_REQUIRE(sd > 0, data_quality, "garch_fit: constant returns");
  const f64 scale = opt.scale_returns ? 1.0 / sd : 1.0;
  std::vector<f64> r(r0.size()); for (usize i = 0; i < r.size(); ++i) r[i] = r0[i] * scale;
  const usize p = garch_n_params(opt.model, opt.dist);
  // starting values (on the scaled series): μ = mean, persistence 0.95, α = 0.05 (GJR γ = 0.05), ν = 8
  std::vector<f64> th0(p); th0[0] = stats::mean(std::span<const f64>(r));
  const f64 v = stats::variance(std::span<const f64>(r), 0);
  const f64 pers = start_variant == 1 ? 0.80 : (start_variant == 2 ? 0.985 : 0.95);
  if (opt.model == GarchModel::egarch) { th0[2] = 0.1; th0[3] = start_variant == 2 ? 0.05 : -0.05; th0[4] = pers; th0[1] = (1.0 - th0[4]) * std::log(v); }
  else if (opt.model == GarchModel::garch) { th0[2] = 0.05; th0[3] = pers - 0.05; th0[1] = v * (1 - pers); }
  else { th0[2] = 0.03; th0[3] = 0.05; th0[4] = pers - 0.055; th0[1] = v * (1 - pers); }
  if (opt.dist == GarchDist::student_t) th0[p - 1] = 8.0;
  auto f_u = [&](const Vec& u) { auto th = garch_from_unconstrained(std::span<const f64>(u.data(), static_cast<usize>(u.size())), opt.model, opt.dist); return garch_nll(r, th, opt.model, opt.dist); };
  auto u0v = garch_to_unconstrained(th0, opt.model, opt.dist);
  Vec x = Eigen::Map<Vec>(u0v.data(), static_cast<Eigen::Index>(p));
  f64 fx = f_u(x); int fevals = 1;
  if (!std::isfinite(fx)) { e.convergence = Convergence{false, 0, fevals, std::numeric_limits<f64>::quiet_NaN(), -1, "invalid starting values"}; return e; }
  auto fcount = [&](const Vec& y) { ++fevals; return f_u(y); };
  Vec g = num_grad(fcount, x, fx); Mat Hinv = Mat::Identity(static_cast<Eigen::Index>(p), static_cast<Eigen::Index>(p));
  int it = 0; bool converged = false; std::string msg;
  for (; it < opt.max_iter; ++it) {
    Vec dir = -Hinv * g; if (dir.dot(g) >= 0) { Hinv.setIdentity(); dir = -g; }
    f64 step = 1.0; f64 fnew = std::numeric_limits<f64>::infinity(); Vec xnew;
    for (int ls = 0; ls < 60; ++ls) { xnew = x + step * dir; fnew = fcount(xnew); if (std::isfinite(fnew) && fnew <= fx + 1e-4 * step * g.dot(dir)) break; step *= 0.5; }
    if (!std::isfinite(fnew) || fnew > fx) { msg = "line search failed"; break; }
    Vec gnew = num_grad(fcount, xnew, fnew);
    if (std::getenv("RAID_GARCH_TRACE")) { std::fprintf(stderr, "it=%d f=%.6f step=%.3g |g|=%.3e x=", it, fnew, step, gnew.norm()); for (Eigen::Index j = 0; j < x.size(); ++j) std::fprintf(stderr, "%.4f ", xnew[j]); std::fprintf(stderr, " g="); for (Eigen::Index j = 0; j < x.size(); ++j) std::fprintf(stderr, "%.3e ", gnew[j]); std::fprintf(stderr, "\n"); }
    const Vec s = xnew - x, yv = gnew - g; const f64 sy = s.dot(yv);
    if (sy > 1e-12) { const Mat I = Mat::Identity(static_cast<Eigen::Index>(p), static_cast<Eigen::Index>(p)); const f64 rho = 1.0 / sy; Hinv = (I - rho * s * yv.transpose()) * Hinv * (I - rho * yv * s.transpose()) + rho * s * s.transpose(); }
    const f64 rel = std::abs(fx - fnew) / std::max(1.0, std::abs(fx));
    x = xnew; fx = fnew; g = gnew;
    if (g.norm() < 1e-5 || (rel < opt.tol && g.norm() < 1e-2) || (s.norm() < 1e-9 && g.norm() < 1e-2)) { converged = true; msg = g.norm() < 1e-5 ? "gradient norm below 1e-5" : (rel < opt.tol ? "relative NLL change below tol (gradient < 1e-2 in transformed space)" : "step below 1e-9 (gradient < 1e-2)"); ++it; break; }
  }
  if (!converged && msg.empty()) msg = "maximum iterations reached";
  auto th = garch_from_unconstrained(std::span<const f64>(x.data(), static_cast<usize>(x.size())), opt.model, opt.dist);
  // Covariance on the natural parameters (scaled series) by numerical derivatives of the NLL
  auto f_th = [&](const Vec& t) { return garch_nll(r, std::span<const f64>(t.data(), static_cast<usize>(t.size())), opt.model, opt.dist); };
  Vec thv = Eigen::Map<Vec>(th.data(), static_cast<Eigen::Index>(p));
  Mat H = num_hessian(f_th, thv);
  // OPG from per-observation scores (central differences of ℓ_t)
  const usize n = r.size(); Mat scores(static_cast<Eigen::Index>(n), static_cast<Eigen::Index>(p));
  { std::vector<f64> lp, lm; for (Eigen::Index j = 0; j < static_cast<Eigen::Index>(p); ++j) { const f64 h = 1e-5 * std::max(1.0, std::abs(thv[j])); Vec a = thv, b = thv; a[j] += h; b[j] -= h;
      bool oka = per_obs_ll(r, std::span<const f64>(a.data(), p), opt.model, opt.dist, lp), okb = per_obs_ll(r, std::span<const f64>(b.data(), p), opt.model, opt.dist, lm);
      for (usize t = 0; t < n; ++t) scores(static_cast<Eigen::Index>(t), j) = (oka && okb) ? (lp[t] - lm[t]) / (2 * h) : std::numeric_limits<f64>::quiet_NaN(); } }
  Mat opg = scores.transpose() * scores;
  InferenceResult inf; Mat cov;
  Eigen::FullPivLU<Mat> lu(H); bool h_ok = lu.isInvertible() && H.allFinite();
  if (opt.cov == "hessian") { cov = h_ok ? Mat(lu.inverse()) : Mat(Mat::Constant(p, p, std::numeric_limits<f64>::quiet_NaN())); inf.covariance_type = "hessian"; }
  else if (opt.cov == "opg") { Eigen::FullPivLU<Mat> lo(opg); cov = lo.isInvertible() ? Mat(lo.inverse()) : Mat(Mat::Constant(p, p, std::numeric_limits<f64>::quiet_NaN())); inf.covariance_type = "opg"; }
  else { cov = h_ok ? Mat(lu.inverse() * opg * lu.inverse()) : Mat(Mat::Constant(p, p, std::numeric_limits<f64>::quiet_NaN())); inf.covariance_type = "qmle_sandwich"; }
  if (!h_ok) e.warnings.push_back("Hessian singular or non-finite at the optimum; standard errors unavailable");
  if (opt.dist == GarchDist::student_t && th.back() > 450.0) e.warnings.push_back("nu at its upper bound (500): the Student-t density is indistinguishable from Gaussian on this sample");
  // Unscale: r0 = r / scale. μ, ω, σ² scale by 1/scale, 1/scale², 1/scale²; α, β, γ, ν unchanged; EGARCH ω → ω + (1−β) log(1/scale²)
  std::vector<f64> th_out = th; Vec jac = Vec::Ones(static_cast<Eigen::Index>(p));
  th_out[0] = th[0] / scale; jac[0] = 1.0 / scale;
  if (opt.model == GarchModel::egarch) { th_out[1] = th[1] + (1.0 - th[4]) * std::log(1.0 / (scale * scale)); }
  else { th_out[1] = th[1] / (scale * scale); jac[1] = 1.0 / (scale * scale); }
  if (opt.model != GarchModel::egarch) cov = jac.asDiagonal() * cov * jac.asDiagonal();
  else { Mat J = Mat::Identity(p, p); J(0, 0) = 1.0 / scale; J(1, 4) = -std::log(1.0 / (scale * scale)); cov = J * cov * J.transpose(); }
  inf.covariance = cov; const f64 nll_unscaled = fx - static_cast<f64>(n) * std::log(scale);  // f_{r0}(x) = scale · f_r(scale·x) ⇒ −Σ log f(r0) = −Σ log f(r) − n log(scale)
  for (Eigen::Index j = 0; j < static_cast<Eigen::Index>(p); ++j) { const f64 se = std::sqrt(std::max(0.0, cov(j, j))); inf.std_errors.push_back(se); inf.t_stats.push_back(se > 0 ? th_out[static_cast<usize>(j)] / se : std::numeric_limits<f64>::quiet_NaN()); inf.p_values.push_back(est::t_pvalue_two_sided(inf.t_stats.back(), 0)); }
  e.inference = inf; e.coefficients = th_out; e.loss = nll_unscaled;
  e.convergence = Convergence{converged, it, fevals, g.norm(), converged ? 0 : 1, msg};
  std::vector<f64> s2; garch_variance_path(r0, th_out, opt.model, opt.dist, s2);
  e.extra = {{"loglik", -nll_unscaled}, {"scale", scale}, {"sigma2_last", s2.empty() ? 0.0 : s2.back()}, {"persistence", opt.model == GarchModel::garch ? th_out[2] + th_out[3] : (opt.model == GarchModel::gjr ? th_out[2] + th_out[4] + 0.5 * th_out[3] : th_out[4])}, {"aic", 2.0 * nll_unscaled + 2.0 * static_cast<f64>(p)}, {"bic", 2.0 * nll_unscaled + std::log(static_cast<f64>(n)) * static_cast<f64>(p)}};
  if (opt.forecast_horizon > 0 && converged) e.extra["forecast_variance"] = garch_forecast(e, r0, opt.forecast_horizon, opt.plan.seed);
  e.runtime_seconds = std::chrono::duration<f64>(std::chrono::steady_clock::now() - t0).count();
  return e;
}

std::vector<f64> garch_forecast(const Estimate& fit, std::span<const f64> r, int H, u64 seed, int n_sim) {
  RAID_REQUIRE(H >= 1, invalid_argument, "horizon >= 1");
  const std::string& m = fit.model; GarchModel model = m.find(".gjr.") != std::string::npos ? GarchModel::gjr : (m.find(".egarch.") != std::string::npos ? GarchModel::egarch : GarchModel::garch);
  GarchDist dist = m.find("student_t") != std::string::npos ? GarchDist::student_t : GarchDist::gaussian;
  const auto& th = fit.coefficients; std::vector<f64> s2; RAID_REQUIRE(garch_variance_path(r, th, model, dist, s2), numerical, "forecast: invalid variance path");
  const f64 mu = th[0], omega = th[1], alpha = th[2]; const f64 gamma = model != GarchModel::garch ? th[3] : 0.0; const f64 beta = th[model != GarchModel::garch ? 4 : 3];
  const f64 eT = r.back() - mu, s2T = s2.back(); std::vector<f64> f(static_cast<usize>(H));
  if (model != GarchModel::egarch) {
    f[0] = omega + (alpha + (eT < 0 ? gamma : 0.0)) * eT * eT + beta * s2T;
    for (int h = 1; h < H; ++h) f[static_cast<usize>(h)] = omega + (alpha + 0.5 * gamma + beta) * f[static_cast<usize>(h - 1)];
    return f;
  }
  // EGARCH: h=1 analytical, h>1 by simulation of z (Gaussian or standardized t)
  const f64 nu = dist == GarchDist::student_t ? th.back() : 0.0;
  const f64 Eabs = dist == GarchDist::student_t ? 2.0 * std::sqrt(nu - 2.0) * std::exp(std::lgamma((nu + 1) / 2) - std::lgamma(nu / 2)) / ((nu - 1) * std::sqrt(std::numbers::pi)) : std::sqrt(2.0 / std::numbers::pi);
  const f64 zT = eT / std::sqrt(s2T); const f64 l1 = omega + alpha * (std::abs(zT) - Eabs) + gamma * zT + beta * std::log(s2T); f[0] = std::exp(l1);
  if (H == 1) return f;
  std::vector<f64> acc(static_cast<usize>(H), 0.0); Rng rng(seed, 7);
  for (int s = 0; s < n_sim; ++s) { f64 ls = l1; acc[0] += std::exp(ls); for (int h = 1; h < H; ++h) { const f64 z = dist == GarchDist::student_t ? rng.student_t(nu) : rng.normal(); ls = omega + alpha * (std::abs(z) - Eabs) + gamma * z + beta * ls; acc[static_cast<usize>(h)] += std::exp(ls); } }
  for (int h = 1; h < H; ++h) f[static_cast<usize>(h)] = acc[static_cast<usize>(h)] / n_sim;
  return f;
}

}  // namespace raid::models