library(here)
library(tidyverse)
library(gt)
library(AER)
source(here("04-topics/rep-chv2011/Rcode/chv2011-data-prep.R"))
source(here("04-topics/rep-chv2011/Rcode/chv2011-mte-core.R"))
source(here("04-topics/rep-chv2011/Rcode/chv2011-gt-quarto.R"))

B_BOOT <- as.integer(Sys.getenv("CHV2011_BOOT", "50"))

ana <- load_chv2011_analysis(trim = TRUE)
df <- ana$data
d <- wage_interaction_terms(df)

ols_iv <- estimate_ols_iv_returns(df)

iv_single <- function(data, inst) {
  dd <- wage_interaction_terms(data)
  fml <- as.formula(paste(
    "wage ~ state +", paste(chv2011_x, collapse = " + "),
    "|", paste(chv2011_x, collapse = " + "), "+", inst
  ))
  coef(summary(ivreg(fml, data = dd)))["state", "Estimate"] / 4
}

iv_all_est <- function(data) {
  dd <- wage_interaction_terms(data)
  fml <- as.formula(paste(
    "wage ~ state +", paste(chv2011_x, collapse = " + "),
    "|", paste(chv2011_x, collapse = " + "),
    "+ pub4 + lwage5_17 + lurate_17 + tuition"
  ))
  coef(summary(ivreg(fml, data = dd)))["state", "Estimate"] / 4
}

methods <- list(
  "OLS" = function(data) estimate_ols_iv_returns(data)$ols,
  "IV (P(Z))" = function(data) estimate_ols_iv_returns(data)$iv,
  "IV (distance to college)" = function(data) iv_single(data, "pub4"),
  "IV (local wage at 17)" = function(data) iv_single(data, "lwage5_17"),
  "IV (local unemployment at 17)" = function(data) iv_single(data, "lurate_17"),
  "IV (tuition at 17)" = function(data) iv_single(data, "tuition"),
  "2SLS (all IVs)" = iv_all_est
)

boot_one <- function(fun) {
  est <- fun(df)
  boots <- replicate(B_BOOT, {
    idx <- sample.int(nrow(df), nrow(df), replace = TRUE)
    fun(df[idx, , drop = FALSE])
  })
  c(est = est, se = sd(boots))
}

table_data <- imap_dfr(methods, function(fun, method) {
  r <- boot_one(fun)
  tibble(
    method = method,
    cell = chv2011_coef_cell(r["est"], r["se"])
  )
})

gt_tbl <- table_data |>
  chv2011_quarto_blank_df() |>
  gt() |>
  tab_header(
    title = "Table A-8",
    subtitle = "OLS and IV estimates of return to a year of college"
  ) |>
  cols_label(method = "Method", cell = "Estimate (SE)")

save(table_data, gt_tbl, file = here("04-topics/rep-chv2011/Rcode/Table_A8.RData"))