#!/usr/bin/env Rscript
#
#  generate.R — Metrika
#
#  Author:  Simon-Pierre Boucher
#  Contact: contact@spboucher.ai
#  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.
#
#  Generates the golden numerical fixtures for MetrikaKit estimator tests
#  (CLAUDE.md §9). Writes a reference dataset (CSV) plus expected values
#  (TSV, key<TAB>value at 17 significant digits) into the Swift test
#  resources at MetrikaKit/Tests/MetrikaKitTests/Fixtures/.
#
#  All expected values are computed from the CSV **after** a write/read
#  round-trip so Swift and R consume bit-identical inputs.
#

out_dir <- file.path(
  dirname(dirname(normalizePath(sub("--file=", "", grep("--file=", commandArgs(FALSE), value = TRUE))))),
  "..", "MetrikaKit", "Tests", "MetrikaKitTests", "Fixtures"
)
out_dir <- normalizePath(out_dir, mustWork = FALSE)
dir.create(out_dir, recursive = TRUE, showWarnings = FALSE)

# ---------------------------------------------------------------- dataset
set.seed(42)
n <- 60
region  <- rep(1:3, length.out = n)
firm_id <- rep(1:12, each = 5)
price   <- round(runif(n, 5, 20), 4)
noise   <- round(rnorm(n, 0, 0.15), 6)
log_rev <- 4 - 0.08 * price + 0.10 * (region == 2) + 0.20 * (region == 3) + noise
revenue <- round(exp(log_rev), 6)

# GLM outcomes — drawn AFTER all draws above so earlier fixture values
# stay bit-identical when regenerating.
purchase <- rbinom(n, 1, plogis(2 - 0.20 * price))
orders   <- rpois(n, exp(1.6 - 0.09 * price))

# Instruments for the 2SLS fixture (correlated with price by
# construction), drawn after everything above.
z1 <- round(price + rnorm(n, 0, 2), 4)
z2 <- round(0.5 * price + rnorm(n, 0, 3), 4)

# Inject missing values to exercise listwise deletion.
price_missing <- price
price_missing[c(7, 23, 41)] <- NA

data <- data.frame(
  revenue = revenue,
  price = price_missing,
  region = region,
  firm_id = firm_id,
  purchase = purchase,
  orders = orders,
  z1 = z1,
  z2 = z2
)
csv_path <- file.path(out_dir, "regression.csv")
write.csv(data, csv_path, row.names = FALSE, quote = FALSE, na = "")

# Round-trip: recompute everything from what was actually written.
data <- read.csv(csv_path)

sink_path <- file.path(out_dir, "expected.tsv")
lines <- character(0)
emit <- function(key, value) {
  lines <<- c(lines, sprintf("%s\t%.17g", key, value))
}

# ------------------------------------------------------------- summarize
x <- data$revenue
emit("sum_revenue_n",    length(x))
emit("sum_revenue_mean", mean(x))
emit("sum_revenue_var",  var(x))
emit("sum_revenue_sd",   sd(x))
emit("sum_revenue_min",  min(x))
emit("sum_revenue_max",  max(x))
m2 <- mean((x - mean(x))^2)
m3 <- mean((x - mean(x))^3)
m4 <- mean((x - mean(x))^4)
emit("sum_revenue_skewness", m3 / m2^1.5)   # Stata definition
emit("sum_revenue_kurtosis", m4 / m2^2)

# ------------------------------------------------------------ regression
# Estimation sample: listwise deletion on price.
complete <- !is.na(data$price)
d <- data[complete, ]
d$log_rev <- log(d$revenue)
n_est <- nrow(d)

X <- cbind(price = d$price, `_cons` = 1)
y <- d$log_rev
k <- ncol(X)

fit <- lm(log_rev ~ price, data = d)
beta <- coef(fit)                      # (Intercept), price
resid <- residuals(fit)
XtXinv <- solve(t(X) %*% X)
df_r <- n_est - k
sigma2 <- sum(resid^2) / df_r

emit("ols_n", n_est)
emit("ols_dropped", sum(!complete))
emit("ols_b_price", beta["price"])
emit("ols_b_cons",  beta["(Intercept)"])
emit("ols_r2",   summary(fit)$r.squared)
emit("ols_r2a",  summary(fit)$adj.r.squared)
emit("ols_rmse", sqrt(sigma2))
emit("ols_F",    summary(fit)$fstatistic[["value"]])

V_classical <- sigma2 * XtXinv
emit("ols_se_classical_price", sqrt(V_classical["price", "price"]))
emit("ols_se_classical_cons",  sqrt(V_classical["_cons", "_cons"]))
t_price <- beta["price"] / sqrt(V_classical["price", "price"])
emit("ols_t_price", t_price)
emit("ols_p_price", 2 * pt(-abs(t_price), df_r))
crit <- qt(0.975, df_r)
emit("ols_ci_lower_price", beta["price"] - crit * sqrt(V_classical["price", "price"]))
emit("ols_ci_upper_price", beta["price"] + crit * sqrt(V_classical["price", "price"]))

# HC0–HC3 sandwich estimators (manual, no packages).
h <- hatvalues(fit)
sandwich_se <- function(w) {
  meat <- t(X * w) %*% X
  V <- XtXinv %*% meat %*% XtXinv
  sqrt(diag(V))
}
se_hc0 <- sandwich_se(resid^2)
se_hc1 <- sandwich_se(resid^2 * n_est / df_r)
se_hc2 <- sandwich_se(resid^2 / (1 - h))
se_hc3 <- sandwich_se(resid^2 / (1 - h)^2)
emit("ols_se_hc0_price", se_hc0["price"])
emit("ols_se_hc0_cons",  se_hc0["_cons"])
emit("ols_se_hc1_price", se_hc1["price"])
emit("ols_se_hc1_cons",  se_hc1["_cons"])
emit("ols_se_hc2_price", se_hc2["price"])
emit("ols_se_hc2_cons",  se_hc2["_cons"])
emit("ols_se_hc3_price", se_hc3["price"])
emit("ols_se_hc3_cons",  se_hc3["_cons"])

# Cluster-robust with Stata regress small-sample factor.
cl <- d$firm_id
u <- rowsum(X * resid, cl)
G <- length(unique(cl))
meat_cl <- (G / (G - 1)) * ((n_est - 1) / df_r) * (t(u) %*% u)
V_cl <- XtXinv %*% meat_cl %*% XtXinv
emit("ols_G", G)
emit("ols_se_cluster_price", sqrt(V_cl["price", "price"]))
emit("ols_se_cluster_cons",  sqrt(V_cl["_cons", "_cons"]))
t_cl <- beta["price"] / sqrt(V_cl["price", "price"])
emit("ols_p_cluster_price", 2 * pt(-abs(t_cl), G - 1))

# Factor-variable regression: log_rev ~ price + i.region (base = 1).
fit2 <- lm(log_rev ~ price + factor(region), data = d)
b2 <- coef(fit2)
emit("ols2_b_price",   b2[["price"]])
emit("ols2_b_region2", b2[["factor(region)2"]])
emit("ols2_b_region3", b2[["factor(region)3"]])
emit("ols2_b_cons",    b2[["(Intercept)"]])
emit("ols2_r2", summary(fit2)$r.squared)

# ------------------------------------------------------------------ GLMs
# Fit tightly (epsilon 1e-12) so R and Swift land on the same optimum to
# well past the 1e-10 test tolerance. Estimation sample = complete cases.
ctrl <- glm.control(epsilon = 1e-12, maxit = 100)

# Expected information (X'WX)^-1 evaluated AT the converged coefficients.
# R's vcov(fit) instead reuses the weights of the last IRLS step (at the
# second-to-last iterate), so it is only ~1e-6-accurate by its own
# stopping rule — not good enough for 1e-10 fixtures.
glm_bread <- function(fit) {
  Xg <- model.matrix(fit)
  eta <- as.vector(Xg %*% coef(fit))
  fam <- fit$family
  mu <- fam$linkinv(eta)
  W <- fam$mu.eta(eta)^2 / fam$variance(mu)
  solve(t(Xg * W) %*% Xg)
}

glm_bread_sandwich <- function(fit, score_scale, cl = NULL) {
  Xg <- model.matrix(fit)
  u <- Xg * score_scale
  bread <- glm_bread(fit)
  if (is.null(cl)) {
    meat <- t(u) %*% u
  } else {
    ug <- rowsum(u, cl)
    G <- nrow(ug)
    meat <- (G / (G - 1)) * (t(ug) %*% ug)
  }
  sqrt(diag(bread %*% meat %*% bread))
}

lg <- glm(purchase ~ price, data = d, family = binomial(), control = ctrl)
V_lg <- glm_bread(lg)
emit("logit_b_price", coef(lg)[["price"]])
emit("logit_b_cons",  coef(lg)[["(Intercept)"]])
emit("logit_se_price", sqrt(V_lg["price", "price"]))
emit("logit_se_cons",  sqrt(V_lg["(Intercept)", "(Intercept)"]))
emit("logit_ll",  as.numeric(logLik(lg)))
lg0 <- glm(purchase ~ 1, data = d, family = binomial(), control = ctrl)
emit("logit_ll0", as.numeric(logLik(lg0)))
emit("logit_chi2", 2 * (as.numeric(logLik(lg)) - as.numeric(logLik(lg0))))
emit("logit_chi2p", pchisq(
  2 * (as.numeric(logLik(lg)) - as.numeric(logLik(lg0))), 1, lower.tail = FALSE
))
emit("logit_se_hc0_price", glm_bread_sandwich(lg, d$purchase - fitted(lg))["price"])
emit("logit_se_cluster_price",
     glm_bread_sandwich(lg, d$purchase - fitted(lg), cl = d$firm_id)["price"])

pr <- glm(purchase ~ price, data = d,
          family = binomial(link = "probit"), control = ctrl)
emit("probit_b_price", coef(pr)[["price"]])
emit("probit_b_cons",  coef(pr)[["(Intercept)"]])
emit("probit_se_price", sqrt(glm_bread(pr)["price", "price"]))
emit("probit_ll", as.numeric(logLik(pr)))

ps <- glm(orders ~ price, data = d, family = poisson(), control = ctrl)
emit("pois_b_price", coef(ps)[["price"]])
emit("pois_b_cons",  coef(ps)[["(Intercept)"]])
emit("pois_se_price", sqrt(glm_bread(ps)["price", "price"]))
emit("pois_ll", as.numeric(logLik(ps)))
emit("pois_se_hc0_price", glm_bread_sandwich(ps, d$orders - fitted(ps))["price"])

# ----------------------------------------------------- fixed effects (FE)
# Within estimator, Stata conventions: demean within firm, add back grand
# means (so _cons is reported), df = N − K − G. Cluster VCE on the panel
# with the G/(G−1) factor and t on G−1 df.
fe_within <- function(v, g) {
  gm <- ave(v, g)             # group means
  v - gm + mean(v)
}
g <- d$firm_id
Nw <- nrow(d)
yt <- fe_within(d$log_rev, g)
xt_p <- fe_within(d$price, g)
Xw <- cbind(price = xt_p, `_cons` = 1)
G <- length(unique(g))
Kw <- 1
df_fe <- Nw - Kw - G

fe_fit <- lm.fit(Xw, yt)
b_fe <- fe_fit$coefficients
res_fe <- fe_fit$residuals
XtXinv_w <- solve(t(Xw) %*% Xw)
sigma2_fe <- sum(res_fe^2) / df_fe
emit("fe_N", Nw)
emit("fe_G", G)
emit("fe_b_price", b_fe[["price"]])
emit("fe_b_cons",  b_fe[["_cons"]])
emit("fe_se_price", sqrt(sigma2_fe * XtXinv_w["price", "price"]))
yd <- d$log_rev - ave(d$log_rev, g)
emit("fe_r2_within", 1 - sum(res_fe^2) / sum(yd^2))
emit("fe_rmse", sqrt(sigma2_fe))

# Cluster on the panel: scores use DEMEANED x (constant column as-is).
xd <- d$price - ave(d$price, g)
u_fe <- rowsum(cbind(xd * res_fe, res_fe), g)
meat_fe <- (G / (G - 1)) * (t(u_fe) %*% u_fe)
V_fe_cl <- XtXinv_w %*% meat_fe %*% XtXinv_w
emit("fe_se_cluster_price", sqrt(V_fe_cl[1, 1]))
t_fe_cl <- b_fe[["price"]] / sqrt(V_fe_cl[1, 1])
emit("fe_p_cluster_price", 2 * pt(-abs(t_fe_cl), G - 1))

# ------------------------------------------------------------- 2SLS (IV)
# ivregress 2sls log_rev (price = z1 z2), Stata `small` convention:
# residuals from the ORIGINAL regressors, sigma2 = u'u/(N-K), t stats.
Xiv <- cbind(price = d$price, `_cons` = 1)
Ziv <- cbind(z1 = d$z1, z2 = d$z2, `_cons` = 1)
PZ <- Ziv %*% solve(t(Ziv) %*% Ziv) %*% t(Ziv)
Xhat <- PZ %*% Xiv
XhXhinv <- solve(t(Xhat) %*% Xhat)
b_iv <- XhXhinv %*% t(Xhat) %*% d$log_rev
u_iv <- d$log_rev - Xiv %*% b_iv
df_iv <- Nw - 2
sigma2_iv <- sum(u_iv^2) / df_iv
V_iv <- sigma2_iv * XhXhinv
emit("iv_N", Nw)
emit("iv_b_price", b_iv[1, 1])
emit("iv_b_cons",  b_iv[2, 1])
emit("iv_se_price", sqrt(V_iv[1, 1]))
emit("iv_se_cons",  sqrt(V_iv[2, 2]))
t_iv <- b_iv[1, 1] / sqrt(V_iv[1, 1])
emit("iv_p_price", 2 * pt(-abs(t_iv), df_iv))
emit("iv_r2", 1 - sum(u_iv^2) / sum((d$log_rev - mean(d$log_rev))^2))
emit("iv_rmse", sqrt(sigma2_iv))

# Robust (HC1 with the small factor N/(N-K)) built on projected X.
meat_iv <- t(Xhat * as.vector(u_iv^2)) %*% Xhat * (Nw / df_iv)
V_iv_r <- XhXhinv %*% meat_iv %*% XhXhinv
emit("iv_se_hc1_price", sqrt(V_iv_r[1, 1]))

# ---------------------------------------------------- margins (AME, delta)
# Average marginal effects with delta-method SEs at the converged
# coefficients, mirroring the engine's formulas exactly.
ame_glm <- function(fit, var, fprime, fsecond) {
  Xg <- model.matrix(fit)
  eta <- as.vector(Xg %*% coef(fit))
  fp <- fprime(eta)
  fs <- fsecond(eta)
  b <- coef(fit)[[var]]
  ame <- b * mean(fp)
  grad <- colMeans(Xg * fs) * b
  grad[[var]] <- grad[[var]] + mean(fp)
  V <- glm_bread(fit)
  se <- sqrt(as.numeric(t(grad) %*% V %*% grad))
  c(ame = ame, se = se)
}

m_lg <- ame_glm(
  lg, "price",
  function(e) plogis(e) * (1 - plogis(e)),
  function(e) plogis(e) * (1 - plogis(e)) * (1 - 2 * plogis(e))
)
emit("margins_logit_price", m_lg[["ame"]])
emit("margins_logit_se",    m_lg[["se"]])

m_ps <- ame_glm(ps, "price", exp, exp)
emit("margins_pois_price", m_ps[["ame"]])
emit("margins_pois_se",    m_ps[["se"]])

# ------------------------------------------------------------ elastic net
# glmnet at fixed lambda, tight threshold. price/z1/z2 are correlated by
# construction, so the l1 penalty must actually choose among them.
if (requireNamespace("glmnet", quietly = TRUE)) {
  Xen <- as.matrix(d[, c("price", "z1", "z2", "orders")])
  yen <- d$log_rev

  fit_lasso <- glmnet::glmnet(
    Xen, yen, alpha = 1, lambda = 0.05,
    thresh = 1e-15, maxit = 1e7, standardize = TRUE
  )
  cl <- as.vector(coef(fit_lasso))
  emit("lasso_b_cons",   cl[1])
  emit("lasso_b_price",  cl[2])
  emit("lasso_b_z1",     cl[3])
  emit("lasso_b_z2",     cl[4])
  emit("lasso_b_orders", cl[5])

  fit_enet <- glmnet::glmnet(
    Xen, yen, alpha = 0.4, lambda = 0.02,
    thresh = 1e-15, maxit = 1e7, standardize = TRUE
  )
  ce <- as.vector(coef(fit_enet))
  emit("enet_b_cons",   ce[1])
  emit("enet_b_price",  ce[2])
  emit("enet_b_z1",     ce[3])
  emit("enet_b_z2",     ce[4])
  emit("enet_b_orders", ce[5])
} else {
  cat("glmnet not installed — skipping elastic-net fixtures\n")
}

# ------------------------------------------------------- gradient boosting
# xgboost exact-greedy at fixed settings, no subsampling — deterministic.
# Complete-case features only (missing-value default-direction search
# differs between implementations).
if (requireNamespace("xgboost", quietly = TRUE)) {
  # Full sample: these variables have no missing values, so the engine's
  # listwise deletion keeps all 60 rows.
  Xgb <- as.matrix(data[, c("z1", "z2", "orders")])
  ygb <- log(data$revenue)
  base <- mean(ygb)
  booster <- xgboost::xgboost(
    x = Xgb, y = ygb,
    nrounds = 20, max_depth = 3, learning_rate = 0.3,
    reg_lambda = 1, reg_alpha = 0, min_child_weight = 1,
    subsample = 1, colsample_bytree = 1,
    tree_method = "exact", base_score = base,
    nthread = 1, verbosity = 0
  )
  pred <- predict(booster, Xgb)
  emit("gbm_rmse", sqrt(mean((ygb - pred)^2)))
  emit("gbm_pred_1",  pred[1])
  emit("gbm_pred_17", pred[17])
  emit("gbm_pred_42", pred[42])
  emit("gbm_pred_57", pred[57])
} else {
  cat("xgboost not installed — skipping boosting fixtures\n")
}

# -------------------------------------------------------------- correlate
emit("corr_rev_price", cor(d$revenue, d$price))
emit("corr_rev_logrev", cor(d$revenue, d$log_rev))

# --------------------------------------------------------- distributions
emit("dist_pchisq_3p8_1", pchisq(3.8, 1))
emit("dist_pchisq_25_4",  pchisq(25, 4))
emit("dist_qnorm_0p975",  qnorm(0.975))
emit("dist_pt_2p5_df10",   pt(2.5, 10))
emit("dist_pt_m1p3_df3",   pt(-1.3, 3))
emit("dist_pt_0p05_df57",  pt(0.05, 57))
emit("dist_pf_3p7_2_30",   pf(3.7, 2, 30))
emit("dist_pf_0p5_5_100",  pf(0.5, 5, 100))
emit("dist_qt_0p975_df12", qt(0.975, 12))
emit("dist_qt_0p995_df4",  qt(0.995, 4))
emit("dist_pnorm_1p64",    pnorm(1.64))

# ------------------------------------------------------------- .dta files
# haven-written fixtures for the native DTA reader: same dataset plus a
# string column, in formats 118 (version 14) and 117 (version 13).
if (requireNamespace("haven", quietly = TRUE)) {
  dta <- data
  dta$firm_name <- paste0("firm_", data$firm_id)
  dta$firm_name[5] <- ""   # string missing
  haven::write_dta(dta, file.path(out_dir, "regression_v118.dta"), version = 14)
  haven::write_dta(dta, file.path(out_dir, "regression_v117.dta"), version = 13)
  cat("wrote .dta fixtures (117, 118)\n")
} else {
  cat("haven not installed — skipping .dta fixtures\n")
}

writeLines(lines, sink_path)
cat("wrote", csv_path, "\n")
cat("wrote", sink_path, "with", length(lines), "expected values\n")