# Shared figure builders for CHV2011 replication.
source(here::here("04-topics/rep-chv2011/Rcode/chv2011-data-prep.R"))
source(here::here("04-topics/rep-chv2011/Rcode/chv2011-mte-core.R"))

chv2011_fig_mte <- function(spec = "baseline", title = NULL) {
  ana <- load_chv2011_analysis(spec = spec, trim = TRUE)
  mte <- estimate_mte_polynomial(ana$data)
  ttl <- if (is.null(title)) paste("MTE curve —", spec) else title
  ggplot2::ggplot(
    tibble::tibble(u = mte$grid, mte = mte$mte / 4),
    ggplot2::aes(u, mte)
  ) +
    ggplot2::geom_line(linewidth = 0.8) +
    ggplot2::labs(x = expression(u[S]), y = "MTE (annualized)", title = ttl) +
    ggplot2::theme_minimal()
}

chv2011_fig_normal_mte <- function(title = "Figure 1 — Normal-model MTE") {
  curve <- load_chv2011_normal_mte_curve()
  ggplot2::ggplot(curve, ggplot2::aes(u, mte)) +
    ggplot2::geom_line(linewidth = 0.8, color = "steelblue") +
    ggplot2::labs(
      x = expression(u[S]),
      y = "MTE (annualized)",
      title = title,
      subtitle = if (file.exists(chv2011_normal_mte_path())) {
        "Author mtexvb.out"
      } else {
        "Proxy from OLS state coefficient"
      }
    ) +
    ggplot2::theme_minimal()
}

chv2011_fig_phats <- function(trim = FALSE) {
  ana <- load_chv2011_analysis(trim = trim)
  ggplot2::ggplot(ana$data, ggplot2::aes(phat)) +
    ggplot2::geom_histogram(bins = 40, fill = "gray70", color = "white") +
    ggplot2::labs(
      title = if (trim) "Propensity score (trimmed support)" else "Propensity score distribution",
      x = "P(Z)",
      y = "Count"
    ) +
    ggplot2::theme_minimal()
}

chv2011_fig_weights <- function(title = "MPRTE weights") {
  g <- chv2011_mte_grid
  tibble::tibble(
    u = rep(g, 3),
    weight = c(
      compute_mprte_weights(g, "z_shift"),
      compute_mprte_weights(g, "p_shift"),
      compute_mprte_weights(g, "p_scale")
    ),
    policy = rep(c("Z+alpha", "P+alpha", "(1+alpha)P"), each = length(g))
  ) |>
    ggplot2::ggplot(ggplot2::aes(u, weight, color = policy)) +
    ggplot2::geom_line(linewidth = 0.8) +
    ggplot2::labs(title = title, x = expression(u[S]), y = "Weight") +
    ggplot2::theme_minimal()
}

chv2011_fig_iv_support <- function() {
  ana <- load_chv2011_analysis(trim = FALSE)
  df <- ana$data
  inst <- df |>
    transmute(
      pub4, lwage5_17, lurate_17, tuition,
      phat
    ) |>
    pivot_longer(-phat, names_to = "instrument", values_to = "z") |>
    mutate(instrument = recode(
      instrument,
      pub4 = "Distance to college",
      lwage5_17 = "Local wage at 17",
      lurate_17 = "Local unemployment at 17",
      tuition = "Tuition at 17"
    ))
  ggplot2::ggplot(inst, ggplot2::aes(z, phat)) +
    ggplot2::geom_point(alpha = 0.15, size = 0.8) +
    ggplot2::facet_wrap(~instrument, scales = "free_x") +
    ggplot2::labs(
      title = "Figure 6 — Propensity score vs instruments",
      x = "Instrument",
      y = "P(Z)"
    ) +
    ggplot2::theme_minimal()
}

chv2011_save_figure <- function(fig_plot, name) {
  save(fig_plot, file = here::here("04-topics/rep-chv2011/Rcode", paste0(name, ".RData")))
}