#include "raid/statistics/covariance.hpp"
#include "raid/core/error.hpp"
#include <cmath>
#include <map>

namespace raid::stats {

std::string to_string(CovType c) {
  switch (c) { case CovType::classical: return "classical"; case CovType::hc0: return "hc0"; case CovType::hc1: return "hc1"; case CovType::hc2: return "hc2"; case CovType::hc3: return "hc3"; case CovType::hac_bartlett: return "hac_bartlett"; case CovType::cluster: return "cluster"; }
  return "?";
}
CovType covtype_from_string(std::string_view s) {
  if (s == "classical" || s == "nonrobust") return CovType::classical;
  if (s == "hc0" || s == "HC0") return CovType::hc0; if (s == "hc1" || s == "HC1") return CovType::hc1;
  if (s == "hc2" || s == "HC2") return CovType::hc2; if (s == "hc3" || s == "HC3") return CovType::hc3;
  if (s == "hac" || s == "hac_bartlett" || s == "newey_west") return CovType::hac_bartlett;
  if (s == "cluster") return CovType::cluster;
  throw RaidError(ErrorCode::invalid_argument, "unknown covariance type '" + std::string(s) + "'");
}
int default_hac_lags(usize n) { return static_cast<int>(std::floor(4.0 * std::pow(static_cast<f64>(n) / 100.0, 2.0 / 9.0))); }

Eigen::MatrixXd robust_covariance(const Eigen::Ref<const Eigen::MatrixXd>& X, const Eigen::Ref<const Eigen::VectorXd>& e,
                                  const Eigen::Ref<const Eigen::MatrixXd>& xtx_inv, const CovOptions& opt, const Eigen::VectorXd& h) {
  const Eigen::Index n = X.rows(), k = X.cols();
  RAID_REQUIRE(e.size() == n, invalid_argument, "robust_covariance: residual length");
  const f64 nd = static_cast<f64>(n), kd = static_cast<f64>(k);
  switch (opt.type) {
    case CovType::classical: {
      RAID_REQUIRE(n > k, insufficient_samples, "classical covariance needs n > k");
      const f64 s2 = e.squaredNorm() / (nd - kd);
      return s2 * xtx_inv;
    }
    case CovType::hc0: case CovType::hc1: case CovType::hc2: case CovType::hc3: {
      Eigen::VectorXd w = e.array().square();
      if (opt.type == CovType::hc1) w *= nd / (nd - kd);
      if (opt.type == CovType::hc2 || opt.type == CovType::hc3) {
        RAID_REQUIRE(h.size() == n, invalid_argument, "HC2/HC3 need leverage");
        for (Eigen::Index i = 0; i < n; ++i) { const f64 d = 1.0 - h[i]; w[i] /= (opt.type == CovType::hc2 ? d : d * d); }
      }
      const Eigen::MatrixXd meat = X.transpose() * w.asDiagonal() * X;
      return xtx_inv * meat * xtx_inv;
    }
    case CovType::hac_bartlett: {
      const int L = opt.hac_lags < 0 ? default_hac_lags(static_cast<usize>(n)) : opt.hac_lags;
      Eigen::MatrixXd S = Eigen::MatrixXd::Zero(k, k);
      Eigen::MatrixXd U = X.array().colwise() * e.array();  // scores n×k
      S = U.transpose() * U;
      for (int l = 1; l <= L; ++l) {
        const f64 wl = 1.0 - static_cast<f64>(l) / (static_cast<f64>(L) + 1.0);
        const Eigen::MatrixXd G = U.bottomRows(n - l).transpose() * U.topRows(n - l);
        S += wl * (G + G.transpose());
      }
      if (opt.hac_small_sample) S *= nd / (nd - kd);
      return xtx_inv * S * xtx_inv;
    }
    case CovType::cluster: {
      RAID_REQUIRE(static_cast<Eigen::Index>(opt.clusters.size()) == n, invalid_argument, "cluster ids length");
      std::map<i64, Eigen::VectorXd> sums;
      for (Eigen::Index i = 0; i < n; ++i) {
        auto it = sums.find(opt.clusters[static_cast<usize>(i)]);
        if (it == sums.end()) it = sums.emplace(opt.clusters[static_cast<usize>(i)], Eigen::VectorXd::Zero(k)).first;
        it->second += X.row(i).transpose() * e[i];
      }
      Eigen::MatrixXd S = Eigen::MatrixXd::Zero(k, k);
      for (auto& [g, s] : sums) S += s * s.transpose();
      const f64 G = static_cast<f64>(sums.size());
      f64 corr = 1.0;
      if (opt.cluster_small_sample) { RAID_REQUIRE(G > 1 && n > k, insufficient_samples, "cluster correction needs G > 1, n > k"); corr = (G / (G - 1.0)) * ((nd - 1.0) / (nd - kd)); }
      return corr * (xtx_inv * S * xtx_inv);
    }
  }
  throw RaidError(ErrorCode::invalid_argument, "unknown CovType");
}

}  // namespace raid::stats