library(here)
library(tidyverse)
library(gt)
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 <- as.integer(Sys.getenv("CHV2011_BOOT", "30"))

specs_6a <- c("baseline", "nodrop", "dropdummies")
specs_6a_labels <- c("Baseline", "No dropouts", "Dropout dummies in wage eq.")

specs_6b <- c("allxinz", "nointall", "notuitnounemp", "notuit", "sls")
specs_6b_labels <- c(
  "All X in Z", "No Z x X interactions", "Cameron-Taber IVs",
  "No tuition IV", "SLS index (author cached)"
)

table6_params <- c("ATE_tilde", "TT_tilde", "TUT_tilde", "MPRTE_Z")

table6_param_labels <- c(
  "ATE_tilde" = "ATE (tilde)",
  "TT_tilde" = "TT (tilde)",
  "TUT_tilde" = "TUT (tilde)",
  "MPRTE_Z" = "MPRTE (|Z gamma - V| < e)"
)

#' Paper Table 6(b) SLS column — author point estimates (no SE; R npindex >24h).
chv2011_table6_sls_paper <- function() {
  c(
    ATE_tilde = 0.0871,
    TT_tilde = 0.2154,
    TUT_tilde = -0.0337,
    MPRTE_Z = 0.0799
  )
}

chv2011_table6_paper_targets <- function() {
  list(
    table_6a = tribble(
      ~parameter, ~Baseline, ~`No dropouts`, ~`Dropout dummies in wage eq.`,
      "ATE_tilde", 0.0815, 0.1246, 0.0995,
      "TT_tilde", 0.2420, 0.2605, 0.2500,
      "TUT_tilde", 0.0135, 0.0274, -0.0388,
      "MPRTE_Z", 0.0802, 0.1104, 0.0988
    ),
    table_6b = tribble(
      ~parameter, ~`All X in Z`, ~`No Z x X interactions`, ~`Cameron-Taber IVs`,
      ~`No tuition IV`, ~`SLS index (author cached)`,
      "ATE_tilde", 0.1409, 0.1208, 0.0851, 0.0626, 0.0871,
      "TT_tilde", 0.2233, 0.2125, 0.2409, 0.2056, 0.2154,
      "TUT_tilde", 0.0135, 0.0350, -0.0570, -0.0682, -0.0337,
      "MPRTE_Z", 0.0802, 0.1156, 0.0821, 0.0591, 0.0799
    )
  )
}

run_col <- function(spec) {
  if (identical(spec, "sls")) {
    est <- chv2011_table6_sls_paper()
    return(tibble(
      parameter = names(est),
      estimate = unname(est),
      se = NA_real_
    ))
  }
  tryCatch({
    r <- run_chv2011_spec(spec, B_boot = B)
    r$semiparam |>
      filter(.data$parameter %in% table6_params) |>
      select(parameter, estimate, se)
  }, error = function(e) {
    tibble(
      parameter = table6_params,
      estimate = NA_real_,
      se = NA_real_
    )
  })
}

format_table6_col <- function(col_df, params = table6_params) {
  map_chr(params, function(p) {
    row <- col_df[col_df$parameter == p, , drop = FALSE]
    if (!nrow(row) || is.na(row$estimate[1])) {
      return(" ")
    }
    if (is.na(row$se[1])) {
      return(chv2011_fmt_num(row$estimate[1]))
    }
    chv2011_coef_cell(row$estimate[1], row$se[1])
  })
}

build_table6 <- function(specs, labels) {
  cols <- map(specs, run_col)
  out <- tibble(parameter = unname(table6_param_labels[table6_params]))
  for (i in seq_along(labels)) {
    out[[labels[i]]] <- format_table6_col(cols[[i]])
  }
  out
}

table_6a <- build_table6(specs_6a, specs_6a_labels)
table_6b <- build_table6(specs_6b, specs_6b_labels)

table_data <- list(
  table_6a = table_6a,
  table_6b = table_6b,
  paper_targets = chv2011_table6_paper_targets(),
  meta = list(
    B = B,
    specs_6a = specs_6a,
    specs_6b = specs_6b,
    sls_source = "author paper Table 6(b); npindex not run (see replication_aer/README.txt)"
  )
)

gt_6a <- table_data$table_6a |>
  chv2011_quarto_blank_df() |>
  gt() |>
  tab_header(
    title = "Table 6(a)",
    subtitle = sprintf(
      "Returns — sensitivity (outcome / sample); bootstrap B = %d (R polynomial MTE core)",
      B
    )
  ) |>
  cols_label(parameter = " ")

gt_6b <- table_data$table_6b |>
  chv2011_quarto_blank_df() |>
  gt() |>
  tab_header(
    title = "Table 6(b)",
    subtitle = sprintf(
      "Returns — sensitivity (choice equation); SLS column cached from paper; B = %d",
      B
    )
  ) |>
  cols_label(parameter = " ")

gt_tbl <- list(gt_6a = gt_6a, gt_6b = gt_6b)
save(table_data, gt_tbl, file = here("04-topics/rep-chv2011/Rcode/Table_6.RData"))