#include "raid/estimators/ols.hpp"
#include "raid/core/error.hpp"
#include <chrono>
#include <cmath>

namespace raid::est {

namespace {
// Two-sided p-value from a t statistic with df degrees of freedom (df<=0 → normal). Regularised incomplete beta via
// Lentz continued fraction (Numerical Recipes), FP64. Validated against scipy.stats.t.sf in tests/reference.
f64 betacf(f64 a, f64 b, f64 x) {
  const int MAXIT = 300; const f64 EPS = 3e-16, FPMIN = 1e-300;
  f64 qab = a + b, qap = a + 1, qam = a - 1, c = 1, d = 1 - qab * x / qap;
  if (std::abs(d) < FPMIN) d = FPMIN; d = 1 / d; f64 h = d;
  for (int m = 1; m <= MAXIT; ++m) {
    const f64 m2 = 2.0 * m; f64 aa = m * (b - m) * x / ((qam + m2) * (a + m2));
    d = 1 + aa * d; if (std::abs(d) < FPMIN) d = FPMIN; c = 1 + aa / c; if (std::abs(c) < FPMIN) c = FPMIN; d = 1 / d; h *= d * c;
    aa = -(a + m) * (qab + m) * x / ((a + m2) * (qap + m2));
    d = 1 + aa * d; if (std::abs(d) < FPMIN) d = FPMIN; c = 1 + aa / c; if (std::abs(c) < FPMIN) c = FPMIN; d = 1 / d;
    const f64 del = d * c; h *= del; if (std::abs(del - 1) < EPS) break;
  }
  return h;
}
f64 betai(f64 a, f64 b, f64 x) {
  if (x <= 0) return 0; if (x >= 1) return 1;
  const f64 bt = std::exp(std::lgamma(a + b) - std::lgamma(a) - std::lgamma(b) + a * std::log(x) + b * std::log(1 - x));
  if (x < (a + 1) / (a + b + 2)) return bt * betacf(a, b, x) / a;
  return 1 - bt * betacf(b, a, 1 - x) / b;
}
}  // namespace

f64 t_pvalue_two_sided(f64 t, f64 df) {
  if (!std::isfinite(t)) return std::numeric_limits<f64>::quiet_NaN();
  if (df <= 0) return std::erfc(std::abs(t) / std::sqrt(2.0));
  return betai(df / 2.0, 0.5, df / (df + t * t));
}

OlsResult ols(const Eigen::Ref<const Eigen::MatrixXd>& X0, const Eigen::Ref<const Eigen::VectorXd>& y0, const OlsOptions& opt) {
  const auto t0 = std::chrono::steady_clock::now();
  const Eigen::Index n = X0.rows(), k = X0.cols();
  RAID_REQUIRE(y0.size() == n, invalid_argument, "ols: X rows != y size");
  RAID_REQUIRE(n > 0 && k > 0, insufficient_samples, "ols: empty design");
  RAID_REQUIRE(X0.allFinite() && y0.allFinite(), data_quality, "ols: non-finite values in X or y (drop or impute explicitly)");
  OlsResult r; r.estimate.model = opt.ridge_lambda > 0 ? "ridge" : (opt.weights.empty() ? "ols" : "wls");
  // Weighted transform: X̃ = √W X, ỹ = √W y
  Eigen::MatrixXd X = X0; Eigen::VectorXd y = y0;
  if (!opt.weights.empty()) {
    RAID_REQUIRE(static_cast<Eigen::Index>(opt.weights.size()) == n, invalid_argument, "wls: weights length");
    for (Eigen::Index i = 0; i < n; ++i) { RAID_REQUIRE(opt.weights[static_cast<usize>(i)] >= 0, invalid_argument, "wls: negative weight"); const f64 s = std::sqrt(opt.weights[static_cast<usize>(i)]); X.row(i) *= s; y[i] *= s; }
  }
  // Ridge via augmentation: append √λ·P rows (P = identity on penalised columns)
  Eigen::MatrixXd Xa = X; Eigen::VectorXd ya = y;
  if (opt.ridge_lambda > 0) {
    Eigen::VectorXd pen = Eigen::VectorXd::Constant(k, std::sqrt(opt.ridge_lambda));
    for (usize u : opt.unpenalized) { RAID_REQUIRE(u < static_cast<usize>(k), invalid_argument, "ridge: bad unpenalized index"); pen[static_cast<Eigen::Index>(u)] = 0; }
    Xa.conservativeResize(n + k, k); ya.conservativeResize(n + k);
    Xa.bottomRows(k) = pen.asDiagonal(); ya.tail(k).setZero();
  }
  Eigen::VectorXd beta(k);
  if (opt.solver == Solver::svd) {
    Eigen::JacobiSVD<Eigen::MatrixXd> svd(Xa, Eigen::ComputeThinU | Eigen::ComputeThinV);
    const auto& s = svd.singularValues();
    r.rank = svd.rank(); r.rcond = s[s.size() - 1] / s[0];
    beta = svd.solve(ya);
    Eigen::VectorXd inv_s2 = (s.array() > s[0] * 1e-14).select(1.0 / s.array().square(), 0.0);
    r.xtx_inv = svd.matrixV() * inv_s2.asDiagonal() * svd.matrixV().transpose();
  } else if (opt.solver == Solver::normal_cholesky) {
    const Eigen::MatrixXd A = Xa.transpose() * Xa; Eigen::LLT<Eigen::MatrixXd> llt(A);
    RAID_REQUIRE(llt.info() == Eigen::Success, numerical, "ols: X'X not positive definite");
    beta = llt.solve(Xa.transpose() * ya); r.xtx_inv = llt.solve(Eigen::MatrixXd::Identity(k, k)); r.rank = k;
    Eigen::JacobiSVD<Eigen::MatrixXd> svd(A); r.rcond = svd.singularValues()[k - 1] / svd.singularValues()[0];
  } else {
    Eigen::ColPivHouseholderQR<Eigen::MatrixXd> qr(Xa);
    qr.setThreshold(1e-12);
    r.rank = qr.rank();
    const Eigen::MatrixXd R = qr.matrixR().topLeftCorner(k, k).template triangularView<Eigen::Upper>();
    r.rcond = std::abs(R(k - 1, k - 1)) / std::abs(R(0, 0));
    if (r.rank < k) {
      r.estimate.warnings.push_back("design is rank deficient (rank " + std::to_string(r.rank) + " < " + std::to_string(k) + "); minimum-norm solution via SVD");
      Eigen::JacobiSVD<Eigen::MatrixXd> svd(Xa, Eigen::ComputeThinU | Eigen::ComputeThinV);
      beta = svd.solve(ya);
      const auto& s = svd.singularValues(); Eigen::VectorXd inv_s2 = (s.array() > s[0] * 1e-14).select(1.0 / s.array().square(), 0.0);
      r.xtx_inv = svd.matrixV() * inv_s2.asDiagonal() * svd.matrixV().transpose();
    } else {
      beta = qr.solve(ya);
      // (X'X)^{-1} = P R^{-1} R^{-T} P'
      const Eigen::MatrixXd Rinv = R.template triangularView<Eigen::Upper>().solve(Eigen::MatrixXd::Identity(k, k));
      r.xtx_inv = qr.colsPermutation() * (Rinv * Rinv.transpose()) * qr.colsPermutation().transpose();
    }
  }
  if (r.rcond < opt.rcond_warn) r.estimate.warnings.push_back("ill-conditioned design: rcond = " + std::to_string(r.rcond));
  r.fitted = X0 * beta;                       // unweighted fitted values
  r.residuals = y0 - r.fitted;
  Eigen::VectorXd e = y - X * beta;           // weighted residuals for inference
  const f64 rss = e.squaredNorm();
  const f64 df_resid = static_cast<f64>(n - r.rank);
  r.sigma2 = df_resid > 0 ? rss / df_resid : std::numeric_limits<f64>::quiet_NaN();
  const f64 ybar = y.mean(); const f64 tss = (y.array() - ybar).square().sum();
  r.r2 = tss > 0 ? 1.0 - rss / tss : std::numeric_limits<f64>::quiet_NaN();
  r.r2_adj = tss > 0 && df_resid > 0 ? 1.0 - (rss / df_resid) / (tss / static_cast<f64>(n - 1)) : std::numeric_limits<f64>::quiet_NaN();
  // leverage h_ii = x_i' (X'X)^{-1} x_i
  r.leverage = (X * r.xtx_inv).cwiseProduct(X).rowwise().sum();
  // inference
  InferenceResult inf;
  stats::CovOptions co = opt.cov;
  inf.covariance = stats::robust_covariance(X, e, r.xtx_inv, co, r.leverage);
  inf.covariance_type = stats::to_string(co.type);
  if (co.type == stats::CovType::hac_bartlett) inf.covariance_type += "(" + std::to_string(co.hac_lags < 0 ? stats::default_hac_lags(static_cast<usize>(n)) : co.hac_lags) + ")";
  inf.std_errors.resize(static_cast<usize>(k)); inf.t_stats.resize(static_cast<usize>(k)); inf.p_values.resize(static_cast<usize>(k));
  const f64 df_t = (co.type == stats::CovType::classical) ? df_resid : df_resid;  // statsmodels uses t(df_resid) for OLS by default
  for (Eigen::Index j = 0; j < k; ++j) {
    const f64 se = std::sqrt(std::max(0.0, inf.covariance(j, j)));
    inf.std_errors[static_cast<usize>(j)] = se; inf.t_stats[static_cast<usize>(j)] = se > 0 ? beta[j] / se : std::numeric_limits<f64>::quiet_NaN();
    inf.p_values[static_cast<usize>(j)] = t_pvalue_two_sided(inf.t_stats[static_cast<usize>(j)], df_t);
  }
  r.estimate.inference = inf;
  r.estimate.coefficients.assign(beta.data(), beta.data() + k);
  r.estimate.names = opt.names; if (r.estimate.names.size() != static_cast<usize>(k)) { r.estimate.names.clear(); for (Eigen::Index j = 0; j < k; ++j) r.estimate.names.push_back("x" + std::to_string(j)); }
  r.estimate.loss = rss; r.estimate.n_obs = static_cast<usize>(n);
  r.estimate.convergence = Convergence{true, 1, 1, 0.0, 0, "direct solve"};
  r.estimate.extra = {{"sigma2", r.sigma2}, {"r2", r.r2}, {"r2_adj", r.r2_adj}, {"rank", r.rank}, {"rcond", r.rcond}, {"df_resid", df_resid}, {"ridge_lambda", opt.ridge_lambda}};
  r.estimate.runtime_seconds = std::chrono::duration<f64>(std::chrono::steady_clock::now() - t0).count();
  return r;
}

SufficientStats accumulate(const Eigen::Ref<const Eigen::MatrixXd>& X, const Eigen::Ref<const Eigen::VectorXd>& y) {
  SufficientStats s; s.xtx = X.transpose() * X; s.xty = X.transpose() * y; s.yty = y.squaredNorm(); s.n = static_cast<usize>(X.rows());
  return s;
}

OlsResult ols_from_stats(const SufficientStats& s, const OlsOptions& opt) {
  const auto t0 = std::chrono::steady_clock::now();
  const Eigen::Index k = s.xtx.rows();
  OlsResult r; r.estimate.model = "ols.sufficient_stats";
  Eigen::MatrixXd A = s.xtx;
  if (opt.ridge_lambda > 0) { for (Eigen::Index j = 0; j < k; ++j) { bool pen = true; for (usize u : opt.unpenalized) if (static_cast<Eigen::Index>(u) == j) pen = false; if (pen) A(j, j) += opt.ridge_lambda; } }
  Eigen::LDLT<Eigen::MatrixXd> ldlt(A);
  Eigen::VectorXd beta;
  Eigen::SelfAdjointEigenSolver<Eigen::MatrixXd> es(A);
  const auto& ev = es.eigenvalues();
  r.rcond = ev[0] / ev[k - 1]; r.rank = (ev.array() > ev[k - 1] * 1e-13).count();
  if (ldlt.info() == Eigen::Success && ldlt.isPositive() && r.rank == k) { beta = ldlt.solve(s.xty); r.xtx_inv = ldlt.solve(Eigen::MatrixXd::Identity(k, k)); }
  else {
    r.estimate.warnings.push_back("X'X singular or indefinite: pseudo-inverse solution");
    Eigen::VectorXd inv = (ev.array() > ev[k - 1] * 1e-13).select(1.0 / ev.array(), 0.0);
    r.xtx_inv = es.eigenvectors() * inv.asDiagonal() * es.eigenvectors().transpose(); beta = r.xtx_inv * s.xty;
  }
  if (r.rcond < opt.rcond_warn) r.estimate.warnings.push_back("ill-conditioned X'X: rcond = " + std::to_string(r.rcond));
  const f64 rss = std::max(0.0, s.yty - 2.0 * beta.dot(s.xty) + beta.dot(s.xtx * beta));
  const f64 df_resid = static_cast<f64>(s.n) - static_cast<f64>(r.rank);
  r.sigma2 = df_resid > 0 ? rss / df_resid : std::numeric_limits<f64>::quiet_NaN();
  InferenceResult inf; inf.covariance = r.sigma2 * r.xtx_inv; inf.covariance_type = "classical";
  for (Eigen::Index j = 0; j < k; ++j) { const f64 se = std::sqrt(std::max(0.0, inf.covariance(j, j))); inf.std_errors.push_back(se); inf.t_stats.push_back(se > 0 ? beta[j] / se : std::numeric_limits<f64>::quiet_NaN()); inf.p_values.push_back(t_pvalue_two_sided(inf.t_stats.back(), df_resid)); }
  r.estimate.inference = inf; r.estimate.coefficients.assign(beta.data(), beta.data() + k);
  r.estimate.names = opt.names; if (r.estimate.names.size() != static_cast<usize>(k)) { r.estimate.names.clear(); for (Eigen::Index j = 0; j < k; ++j) r.estimate.names.push_back("x" + std::to_string(j)); }
  r.estimate.loss = rss; r.estimate.n_obs = s.n; r.estimate.convergence = Convergence{true, 1, 1, 0.0, 0, "sufficient statistics"};
  r.estimate.extra = {{"sigma2", r.sigma2}, {"rank", r.rank}, {"rcond", r.rcond}, {"df_resid", df_resid}};
  r.estimate.runtime_seconds = std::chrono::duration<f64>(std::chrono::steady_clock::now() - t0).count();
  return r;
}

RollingOlsResult rolling_ols(const Eigen::Ref<const Eigen::MatrixXd>& X, const Eigen::Ref<const Eigen::VectorXd>& y, usize window, usize step) {
  const usize n = static_cast<usize>(X.rows()), k = static_cast<usize>(X.cols());
  RAID_REQUIRE(window >= k + 1 && window <= n, insufficient_samples, "rolling_ols: window must satisfy k+1 <= window <= n");
  RAID_REQUIRE(step >= 1, invalid_argument, "rolling_ols: step >= 1");
  const usize nw = (n - window) / step + 1;
  RollingOlsResult r; r.window = window; r.coefficients.resize(static_cast<Eigen::Index>(k), static_cast<Eigen::Index>(nw)); r.rss.resize(nw); r.ranks.resize(nw);
  for (usize j = 0; j < nw; ++j) {
    const usize b = j * step;
    Eigen::ColPivHouseholderQR<Eigen::MatrixXd> qr(X.middleRows(static_cast<Eigen::Index>(b), static_cast<Eigen::Index>(window)));
    qr.setThreshold(1e-12);
    const Eigen::VectorXd beta = qr.solve(y.segment(static_cast<Eigen::Index>(b), static_cast<Eigen::Index>(window)));
    r.coefficients.col(static_cast<Eigen::Index>(j)) = beta;
    r.rss[j] = (y.segment(static_cast<Eigen::Index>(b), static_cast<Eigen::Index>(window)) - X.middleRows(static_cast<Eigen::Index>(b), static_cast<Eigen::Index>(window)) * beta).squaredNorm();
    r.ranks[j] = qr.rank();
  }
  return r;
}

}  // namespace raid::est