// OLS / WLS / ridge with stable linear algebra (Householder QR by default, SVD for rank-deficient designs).
#pragma once
#include <vector>
#include <span>
#include <optional>
#include <Eigen/Dense>
#include "raid/core/types.hpp"
#include "raid/core/estimate.hpp"
#include "raid/statistics/covariance.hpp"

namespace raid::est {

enum class Solver { qr, svd, normal_cholesky /* reference only: ill-conditioned designs will be flagged */ };

struct OlsOptions {
  stats::CovOptions cov{};
  Solver solver{Solver::qr};
  f64 ridge_lambda{0.0};            // L2 penalty on all coefficients except those in `unpenalized` (e.g. the intercept)
  std::vector<usize> unpenalized{}; // column indices not penalised
  std::span<const f64> weights{};   // WLS weights w_i (variance ∝ 1/w_i); empty = OLS
  f64 rcond_warn{1e-10};            // warn when the design's reciprocal condition number is below this
  std::vector<std::string> names{};
};

struct OlsResult {
  Estimate estimate;                // coefficients, inference, loss = RSS
  Eigen::VectorXd residuals;
  Eigen::VectorXd fitted;
  Eigen::MatrixXd xtx_inv;          // (X'WX + λP)^{-1}
  f64 sigma2{0};                    // RSS/(n−k)
  f64 r2{0}, r2_adj{0};
  Eigen::Index rank{0};
  f64 rcond{0};
  Eigen::VectorXd leverage;         // h_ii
};

OlsResult ols(const Eigen::Ref<const Eigen::MatrixXd>& X, const Eigen::Ref<const Eigen::VectorXd>& y, const OlsOptions& opt = {});

/// Sufficient statistics of a (possibly partial) regression block: X'X, X'y, y'y, n. Batched/GPU accumulation
/// produces these and the solve happens in FP64 on the CPU.
struct SufficientStats {
  Eigen::MatrixXd xtx; Eigen::VectorXd xty; f64 yty{0}; usize n{0};
  void merge(const SufficientStats& o) { if (n == 0) { *this = o; return; } xtx += o.xtx; xty += o.xty; yty += o.yty; n += o.n; }
};
SufficientStats accumulate(const Eigen::Ref<const Eigen::MatrixXd>& X, const Eigen::Ref<const Eigen::VectorXd>& y);
/// Solve from sufficient statistics (Cholesky with symmetric pivoting fallback to eigen-decomposition); classical SE only.
OlsResult ols_from_stats(const SufficientStats& s, const OlsOptions& opt = {});

/// Rolling OLS: windows [i−w+1, i] for i ≥ w−1; returns k × n_windows coefficients (column j = window ending at index w−1+j).
struct RollingOlsResult { Eigen::MatrixXd coefficients; std::vector<f64> rss; std::vector<Eigen::Index> ranks; usize window; };
RollingOlsResult rolling_ols(const Eigen::Ref<const Eigen::MatrixXd>& X, const Eigen::Ref<const Eigen::VectorXd>& y, usize window, usize step = 1);

}  // namespace raid::est

namespace raid::est {
/// Two-sided p-value of a t statistic (df ≤ 0 → standard normal). FP64, validated against scipy in tests/reference.
f64 t_pvalue_two_sided(f64 t, f64 df);
}