#include "raid/models/garch.hpp"
#include "raid/core/error.hpp"
#include "raid/statistics/reductions.hpp"
#include <cmath>
#include <limits>
#include <numbers>

namespace raid::models {

constexpr f64 kNuMin = 2.05, kNuMax = 500.0;  // same bounds as the arch package (StudentsT distribution bounds)

std::string to_string(GarchModel m) { switch (m) { case GarchModel::garch: return "garch"; case GarchModel::gjr: return "gjr"; case GarchModel::egarch: return "egarch"; } return "?"; }
std::string to_string(GarchDist d) { return d == GarchDist::gaussian ? "gaussian" : "student_t"; }

usize garch_n_params(GarchModel m, GarchDist d) { return 4 + (m != GarchModel::garch ? 1 : 0) + (d == GarchDist::student_t ? 1 : 0); }
std::vector<std::string> garch_param_names(GarchModel m, GarchDist d) {
  std::vector<std::string> n = {"mu", "omega", "alpha"}; if (m != GarchModel::garch) n.push_back("gamma"); n.push_back("beta"); if (d == GarchDist::student_t) n.push_back("nu"); return n;
}

bool garch_params_valid(std::span<const f64> th, GarchModel m, GarchDist d) {
  if (th.size() != garch_n_params(m, d)) return false;
  for (f64 v : th) if (!std::isfinite(v)) return false;
  const f64 omega = th[1], alpha = th[2];
  const f64 gamma = m != GarchModel::garch ? th[3] : 0.0;
  const f64 beta = th[m != GarchModel::garch ? 4 : 3];
  if (d == GarchDist::student_t && th.back() <= 2.0) return false;
  switch (m) {
    case GarchModel::garch: return omega > 0 && alpha >= 0 && beta >= 0 && alpha + beta < 1.0;
    case GarchModel::gjr: return omega > 0 && alpha >= 0 && beta >= 0 && alpha + gamma >= 0 && alpha + beta + 0.5 * gamma < 1.0;
    case GarchModel::egarch: return std::abs(beta) < 1.0;
  }
  return false;
}

bool garch_variance_path(std::span<const f64> r, std::span<const f64> th, GarchModel m, GarchDist d, std::vector<f64>& s2) {
  const usize n = r.size(); s2.assign(n, std::numeric_limits<f64>::quiet_NaN());
  if (n < 2) return false;
  const f64 mu = th[0], omega = th[1], alpha = th[2];
  const f64 gamma = m != GarchModel::garch ? th[3] : 0.0;
  const f64 beta = th[m != GarchModel::garch ? 4 : 3];
  // initial variance: sample variance of ε (ddof=0), as arch's default for the first observation uses backcast... we use
  // the plain sample variance — documented in the registry (tests pin this convention against arch with backcast given).
  f64 m1 = 0; for (f64 v : r) m1 += (v - mu); m1 /= static_cast<f64>(n);
  f64 var0 = 0; for (f64 v : r) var0 += (v - mu - m1) * (v - mu - m1); var0 /= static_cast<f64>(n);
  if (!(var0 > 0)) return false;
  if (m == GarchModel::egarch) {
    const f64 Eabs = (d == GarchDist::student_t) ? [&] { const f64 nu = th.back(); return 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);
    f64 ls = std::log(var0);
    s2[0] = var0;
    for (usize t = 1; t < n; ++t) {
      const f64 e = r[t - 1] - mu, z = e / std::sqrt(s2[t - 1]);
      ls = omega + alpha * (std::abs(z) - Eabs) + gamma * z + beta * std::log(s2[t - 1]);
      s2[t] = std::exp(ls);
      if (!std::isfinite(s2[t]) || s2[t] <= 0) return false;
    }
    return true;
  }
  s2[0] = var0;
  for (usize t = 1; t < n; ++t) {
    const f64 e = r[t - 1] - mu; const f64 e2 = e * e;
    s2[t] = omega + (alpha + (e < 0 ? gamma : 0.0)) * e2 + beta * s2[t - 1];
    if (!std::isfinite(s2[t]) || s2[t] <= 0) return false;
  }
  return true;
}

f64 garch_nll(std::span<const f64> r, std::span<const f64> th, GarchModel m, GarchDist d) {
  if (!garch_params_valid(th, m, d)) return std::numeric_limits<f64>::infinity();
  std::vector<f64> s2;
  if (!garch_variance_path(r, th, m, d, s2)) return std::numeric_limits<f64>::infinity();
  const f64 mu = th[0]; const usize n = r.size();
  f64 nll = 0, c = 0;
  auto add = [&](f64 v) { const f64 t = nll + v; if (std::abs(nll) >= std::abs(v)) c += (nll - t) + v; else c += (v - t) + nll; nll = t; };
  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; add(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; add(-(k - 0.5 * std::log(s2[t]) - 0.5 * (nu + 1) * std::log1p(e * e / ((nu - 2) * s2[t])))); }
  }
  return nll + c;
}

f64 garch_nll(std::span<const f64> r, std::span<const f64> th, GarchModel m, f64 nu_fixed) {
  if (nu_fixed > 2.0) { std::vector<f64> t(th.begin(), th.end()); t.push_back(nu_fixed); return garch_nll(r, t, m, GarchDist::student_t); }
  return garch_nll(r, th, m, GarchDist::gaussian);
}

// Transforms: u = [μ, log ω, logit-type for (α, β[, γ]), log(ν − 2)]
std::vector<f64> garch_to_unconstrained(std::span<const f64> th, GarchModel m, GarchDist d) {
  std::vector<f64> u(th.begin(), th.end());
  u[1] = m == GarchModel::egarch ? th[1] : std::log(th[1]);
  const bool g = m != GarchModel::garch; const usize ib = g ? 4 : 3;
  if (m == GarchModel::egarch) { u[2] = th[2]; u[3] = th[3]; u[4] = std::atanh(std::clamp(th[4], -0.999999, 0.999999)); }
  else {
    // persistence p = α + β (+γ/2) ∈ (0,1): u_p = logit(p); share s = α/p ∈ (0,1): u_s = logit(s); γ free for GJR via α+γ ≥ 0 → γ = exp(u_g) − α
    const f64 gamma = g ? th[3] : 0.0; const f64 p = th[2] + th[ib] + 0.5 * gamma; const f64 s = p > 0 ? th[2] / p : 0.5;
    u[2] = std::log(std::clamp(p, 1e-12, 1 - 1e-12) / (1 - std::clamp(p, 1e-12, 1 - 1e-12)));
    u[ib] = std::log(std::clamp(s, 1e-12, 1 - 1e-12) / (1 - std::clamp(s, 1e-12, 1 - 1e-12)));
    if (g) u[3] = std::log(std::max(th[2] + gamma, 1e-12));
  }
  if (d == GarchDist::student_t) { const f64 s = std::clamp((th.back() - kNuMin) / (kNuMax - kNuMin), 1e-9, 1 - 1e-9); u.back() = std::log(s / (1 - s)); }
  return u;
}

std::vector<f64> garch_from_unconstrained(std::span<const f64> u, GarchModel m, GarchDist d) {
  std::vector<f64> th(u.begin(), u.end());
  th[1] = m == GarchModel::egarch ? u[1] : std::exp(u[1]);
  const bool g = m != GarchModel::garch;
  auto sig = [](f64 x) { return 1.0 / (1.0 + std::exp(-x)); };
  if (m == GarchModel::egarch) { th[2] = u[2]; th[3] = u[3]; th[4] = std::tanh(u[4]); }
  else if (!g) { const f64 p = sig(u[2]) * (1 - 1e-9), s = sig(u[3]); th[2] = s * p; th[3] = (1 - s) * p; }
  else {
    // GJR: p = α + β + γ/2, s = α/p → α = s p; ag = α + γ = exp(u_g) → γ = ag − α; β = p − α − γ/2 (must be ≥ 0: clamp via s)
    const f64 p = sig(u[2]) * (1 - 1e-9), s = sig(u[4]); const f64 alpha = s * p; const f64 ag = std::exp(u[3]); const f64 gamma = ag - alpha;
    f64 beta = p - alpha - 0.5 * gamma; if (beta < 0) { beta = 0; }
    th[2] = alpha; th[3] = gamma; th[4] = beta;
  }
  if (d == GarchDist::student_t) th.back() = kNuMin + (kNuMax - kNuMin) * sig(u.back());
  return th;
}

}  // namespace raid::models