#include <catch2/catch_test_macros.hpp>
#include <catch2/matchers/catch_matchers_floating_point.hpp>
#include "raid/estimators/ols.hpp"
#include "raid/core/rng.hpp"
using namespace raid; using Catch::Matchers::WithinAbs; using Catch::Matchers::WithinRel;

static void make_design(Eigen::MatrixXd& X, Eigen::VectorXd& y, int n = 500, u64 seed = 11) {
  Rng r(seed); X.resize(n, 3); y.resize(n);
  for (int i = 0; i < n; ++i) { X(i, 0) = 1; X(i, 1) = r.normal(); X(i, 2) = r.normal() * 2 + 1; y[i] = 0.5 + 1.5 * X(i, 1) - 0.7 * X(i, 2) + r.normal() * (1 + 0.5 * std::abs(X(i, 1))); }
}
TEST_CASE("OLS recovers coefficients; QR, SVD and sufficient statistics agree", "[ols]") {
  Eigen::MatrixXd X; Eigen::VectorXd y; make_design(X, y);
  auto a = est::ols(X, y); est::OlsOptions o2; o2.solver = est::Solver::svd; auto b = est::ols(X, y, o2);
  auto c = est::ols_from_stats(est::accumulate(X, y));
  for (int j = 0; j < 3; ++j) { CHECK_THAT(a.estimate.coefficients[j], WithinRel(b.estimate.coefficients[j], 1e-10)); CHECK_THAT(a.estimate.coefficients[j], WithinRel(c.estimate.coefficients[j], 1e-9)); CHECK_THAT(a.estimate.inference->std_errors[j], WithinRel(c.estimate.inference->std_errors[j], 1e-8)); }
  CHECK_THAT(a.estimate.coefficients[1], WithinAbs(1.5, 0.15)); CHECK(a.estimate.ok()); CHECK(a.rank == 3);
  CHECK_THAT(a.estimate.loss, WithinRel(a.residuals.squaredNorm(), 1e-12));
}
TEST_CASE("Robust covariances are PSD and HC1 = HC0 * n/(n-k)", "[ols]") {
  Eigen::MatrixXd X; Eigen::VectorXd y; make_design(X, y);
  est::OlsOptions o; o.cov.type = stats::CovType::hc0; auto h0 = est::ols(X, y, o); o.cov.type = stats::CovType::hc1; auto h1 = est::ols(X, y, o);
  o.cov.type = stats::CovType::hc3; auto h3 = est::ols(X, y, o); o.cov.type = stats::CovType::hac_bartlett; o.cov.hac_lags = 5; auto hac = est::ols(X, y, o);
  const double f = 500.0 / 497.0;
  for (int j = 0; j < 3; ++j) { CHECK_THAT(h1.estimate.inference->covariance(j, j), WithinRel(f * h0.estimate.inference->covariance(j, j), 1e-12)); CHECK(h3.estimate.inference->std_errors[j] >= h0.estimate.inference->std_errors[j] * 0.99); CHECK(hac.estimate.inference->std_errors[j] > 0); }
  CHECK(hac.estimate.inference->covariance_type == "hac_bartlett(5)");
  std::vector<i64> cl(500); for (int i = 0; i < 500; ++i) cl[i] = i / 25; o.cov.type = stats::CovType::cluster; o.cov.clusters = cl; auto c = est::ols(X, y, o);
  Eigen::SelfAdjointEigenSolver<Eigen::MatrixXd> es(c.estimate.inference->covariance); CHECK(es.eigenvalues().minCoeff() >= -1e-12);
}
TEST_CASE("Rank deficiency and ridge", "[ols]") {
  Eigen::MatrixXd X; Eigen::VectorXd y; make_design(X, y); Eigen::MatrixXd Xd(X.rows(), 4); Xd << X, X.col(1) * 2.0;
  auto r = est::ols(Xd, y); CHECK(r.rank == 3); CHECK_FALSE(r.estimate.warnings.empty());
  est::OlsOptions o; o.ridge_lambda = 10.0; o.unpenalized = {0}; auto rr = est::ols(X, y, o); auto ols = est::ols(X, y);
  CHECK(std::abs(rr.estimate.coefficients[1]) < std::abs(ols.estimate.coefficients[1])); CHECK(rr.estimate.model == "ridge");
}
TEST_CASE("Rolling OLS matches window-by-window OLS", "[ols]") {
  Eigen::MatrixXd X; Eigen::VectorXd y; make_design(X, y, 300);
  auto roll = est::rolling_ols(X, y, 100, 50); CHECK(roll.coefficients.cols() == 5);
  auto w = est::ols(X.middleRows(50, 100), y.segment(50, 100)); for (int j = 0; j < 3; ++j) CHECK_THAT(roll.coefficients(j, 1), WithinRel(w.estimate.coefficients[j], 1e-10));
}
TEST_CASE("t p-values match known values", "[ols]") {
  CHECK_THAT(est::t_pvalue_two_sided(1.96, 0), WithinAbs(0.04999579, 1e-7));
  CHECK_THAT(est::t_pvalue_two_sided(2.0, 10), WithinAbs(0.07338803, 1e-7));   // scipy.stats.t.sf(2,10)*2
  CHECK_THAT(est::t_pvalue_two_sided(-3.0, 30), WithinAbs(0.00538996, 1e-7));
}