adsy_app.R

library(rlang)
library(dplyr)
library(tidyr)
library(glmmTMB)
library(robustbase)
library(multcomp)
library(optmatch)
library(propertee)

widen_eb_preds <- function(eb_preds, sgs, yr) {
  eb_preds <- tidyr::pivot_wider(eb_preds[eb_preds$sg %in% sgs,],
                                 id_cols = "tx_schoolid",
                                 values_from = setdiff(colnames(eb_preds), c("tx_schoolid", "sg", "lang")),
                                 names_from = c("sg", "lang"))
  colnames(eb_preds) <- c("tx_schoolid",
                          vapply(strsplit(colnames(eb_preds)[2:ncol(eb_preds)], "_"),
                                 function(v) paste(v[2], v[3], "rs", v[1], v[4], yr, sep = "_"),
                                 character(1L)))
  eb_preds
}

wt.avg_scrs_by_lang <- function(scr_colnm, df, yr) {
  d_col <- gsub("_rs_", "_d_", scr_colnm)
  e_scr_col <- gsub(paste0("_", yr), paste("", "e", yr, sep = "_"), scr_colnm)
  s_scr_col <- gsub(paste0("_", yr), paste("", "s", yr, sep = "_"), scr_colnm)
  e_d_col <- gsub("_rs_", "_d_", e_scr_col)
  s_d_col <- gsub("_rs_", "_d_", s_scr_col)
  dplyr::transmute(
    df,
    !!dplyr::sym(d_col) := tidyr::replace_na(!!dplyr::sym(e_d_col), 0) + tidyr::replace_na(!!dplyr::sym(s_d_col), 0),
    # !!dplyr::sym(scr_colnm) := dplyr::case_when(
    #   !is.na(!!dplyr::sym(e_scr_col)) & !is.na(!!dplyr::sym(s_scr_col)) ~
    #     tidyr::replace_na((!!dplyr::sym(e_d_col) * !!dplyr::sym(e_scr_col) + !!dplyr::sym(s_d_col) * !!dplyr::sym(s_scr_col)) /
    #                         !!dplyr::sym(d_col), 0),
    #   !is.na(!!dplyr::sym(e_scr_col)) ~ !!dplyr::sym(e_scr_col),
    #   !is.na(!!dplyr::sym(s_scr_col)) ~ !!dplyr::sym(s_scr_col),
    #   TRUE ~ 0
    # )
    !!dplyr::sym(scr_colnm) := dplyr::case_when(
      (is.na(!!dplyr::sym(e_d_col)) | (!is.na(!!dplyr::sym(e_d_col)) & !!dplyr::sym(e_d_col) == 0)) &
        (is.na(!!dplyr::sym(s_d_col)) | (!is.na(!!dplyr::sym(s_d_col)) & !!dplyr::sym(s_d_col) == 0)) ~ 0,
      !is.na(!!dplyr::sym(e_d_col)) & !!dplyr::sym(e_d_col) >= 5 &
        !is.na(!!dplyr::sym(s_d_col)) & !!dplyr::sym(s_d_col) >= 5 ~
        (!!dplyr::sym(e_d_col) * !!dplyr::sym(e_scr_col) + !!dplyr::sym(s_d_col) * !!dplyr::sym(s_scr_col)) /
        !!dplyr::sym(d_col),
      !is.na(!!dplyr::sym(e_d_col)) & !!dplyr::sym(e_d_col) >= 5 ~ !!dplyr::sym(e_scr_col),
      !is.na(!!dplyr::sym(s_d_col)) & !!dplyr::sym(s_d_col) >= 5 ~ !!dplyr::sym(s_scr_col),
      TRUE ~ NA_real_
    )
  )
}

wt.avg_eb_preds_by_lang <- function(scr_colnm, df, yr) {
  d_col <- gsub("_rs_", "_d_", scr_colnm)
  e_scr_col <- gsub(paste0("_", yr), paste("", "e", yr, sep = "_"), scr_colnm)
  s_scr_col <- gsub(paste0("_", yr), paste("", "s", yr, sep = "_"), scr_colnm)
  e_d_col <- gsub("_rs_", "_d_", e_scr_col)
  s_d_col <- gsub("_rs_", "_d_", s_scr_col)
  dplyr::transmute(
    df,
    !!dplyr::sym(scr_colnm) := dplyr::case_when(
      !is.na(!!dplyr::sym(e_d_col)) & !!dplyr::sym(e_d_col) > 0 &
        !is.na(!!dplyr::sym(s_d_col)) & !!dplyr::sym(s_d_col) > 0 ~
        (!!dplyr::sym(e_d_col) * !!dplyr::sym(e_scr_col) +
           !!dplyr::sym(s_d_col) * !!dplyr::sym(s_scr_col)) / !!dplyr::sym(d_col),
      !is.na(!!dplyr::sym(e_d_col)) & !!dplyr::sym(e_d_col) > 0 ~ !!dplyr::sym(e_scr_col),
      !is.na(!!dplyr::sym(s_d_col)) & !!dplyr::sym(e_d_col) > 0 ~ !!dplyr::sym(s_scr_col),
      TRUE ~ 0
    )
  )
}

#' The measurement error variance of the subgroup average aggregated by language is
#' (d_eng * CSEM_eng^2 + d_spa * CSEM_spa^2) / (d_eng+d_spa)^2
wt.avg_meas_error_by_lang <- function(csem_colnm, df, yr) {
  d_colnm <- gsub("_sig_", "_d_", csem_colnm)
  grade_str <- substr(csem_colnm, regexpr("_g\\d{1}", csem_colnm), regexpr("_g\\d{1}", csem_colnm) + 2)
  e_d_col <- gsub(grade_str, paste0(grade_str, "_e"), d_colnm)
  s_d_col <- gsub(grade_str, paste0(grade_str, "_s"), d_colnm)
  e_csem_col <- gsub(grade_str, paste0(grade_str, "_e"), csem_colnm)
  s_csem_col <- gsub(grade_str, paste0(grade_str, "_s"), csem_colnm)

  dplyr::transmute(
    df,
    !!dplyr::sym(csem_colnm) := dplyr::case_when(
      !is.na(!!dplyr::sym(e_d_col)) & !!dplyr::sym(e_d_col) > 0 &
        !is.na(!!dplyr::sym(s_d_col)) & !!dplyr::sym(s_d_col) > 0 ~
        sqrt((!!dplyr::sym(e_d_col) * (!!dplyr::sym(e_csem_col))^2 + !!dplyr::sym(s_d_col) * (!!dplyr::sym(s_csem_col))^2) /
               (!!dplyr::sym(d_colnm))^2),
      !is.na(!!dplyr::sym(e_d_col)) & !!dplyr::sym(e_d_col) > 0 ~ sqrt((!!dplyr::sym(e_csem_col))^2/!!dplyr::sym(e_d_col)),
      !is.na(!!dplyr::sym(s_d_col)) & !!dplyr::sym(s_d_col) > 0 ~ sqrt((!!dplyr::sym(s_csem_col))^2/!!dplyr::sym(s_d_col)),
      TRUE ~ 0
    )
  )
}

get_one_csem <- function(scr, subj, lang, year, grade) {
  if (!exists("CSEM", envir = globalenv())) stop("Run `source('adsy_app_constants.R')` to retrieve CSEM table")
  if (is.na(scr)) return(NA_real_)
    
  ix <- findInterval(
    scr,
    CSEM$X[CSEM$subj == subj & CSEM$lang == lang &
             CSEM$year == year & CSEM$grade == grade & !is.na(CSEM$X)]
  )
  CSEM$csem[CSEM$subj == subj & CSEM$lang == lang &
              CSEM$year == year & CSEM$grade == grade & !is.na(CSEM$X)][ix]
}

get_vec_csem <- function(scr_col, lang_col, subj, grade, year) {
  if (length(scr_col) != length(lang_col)) stop("scr_col and lang_col should have the same length")
  vapply(
    seq_along(scr_col),
    function(ix) get_one_csem(scr = scr_col[ix], subj = subj
                              , lang = ifelse(subj == "m", "e", lang_col[ix]) # CSEM values are the same in english and spanish for math
                              , year = year, grade = grade),
    numeric(1L)
  )
}


fit_memod <- function(sbgrp_df, fixef_cols, re_specs, grade, subj, year, sg = NULL, lang = NULL, ids_subset = NULL, ...) {
  if (length(setdiff(c("rs", "join_yr", "lang", "sg", "grade", "subj", "tx_schoolid", "d"), colnames(sbgrp_df))) > 0) stop(
    "sbgrp_df must have rs, join_yr, lang, grade, sg, subj, tx_schoolid, and d columns"
  )
  if (!is.null(sg)) sbgrp_df <- sbgrp_df[sbgrp_df$sg == sg,]
  if (!is.null(lang)) sbgrp_df <- sbgrp_df[sbgrp_df$lang == lang,]
  if (!is.null(year)) sbgrp_df <- sbgrp_df[sbgrp_df$join_yr == year,]
  if (!is.null(ids_subset)) sbgrp_df <- sbgrp_df[sbgrp_df$tx_schoolid %in% ids_subset,]
  fit_df <- sbgrp_df[sbgrp_df$subj == subj & sbgrp_df$grade == grade,]
  
  form <- reformulate(c(fixef_cols, re_specs), response = "rs")
  ctr <- do.call(glmmTMB::glmmTMBControl, list(...))

  init_sig_vals <- get_vec_csem(fit_df$rs, fit_df$lang, subj, grade, year)
  first_fit <- glmmTMB::glmmTMB(form, fit_df, weights = d, family = gaussian()
                                , dispformula = ~0-offset(log(init_sig_vals/sqrt(d)))
                                , control = ctr)
  Xpred <- predict(first_fit, type = "response")
  use_sig_vals <- get_vec_csem(Xpred, fit_df[row.names(first_fit$frame), "lang"], subj, grade, year)
  # mapply(
  #   function(rs, lang, subj, grade, year) get_csem(scr = rs, subj = subj
  #                               , lang = ifelse(subj == "m", "e", lang) # CSEM values are the same in english and spanish for math
  #                               , year = year, grade = grade),
  #   Xpred,
  #   fit_df[row.names(first_fit$frame), "lang"],
  #   MoreArgs = list(subj = subj, grade = grade, year = year)
  # )
  glmmTMB::glmmTMB(
    form, fit_df[row.names(first_fit$frame),], weights = d, family = gaussian()
    , dispformula = ~0-offset(log(use_sig_vals/sqrt(d))), control = ctr
  )
}

# fit_memods <- function(sbgrp_df, W_col, grade_col, Z_cols, id_col, sbgrp_col, subj, lang = NULL, year = NULL, ...) {
#   if (!is.null(year)) sbgrp_df <- sbgrp_df[sbgrp_df$join_yr == year,]
#   if (!is.null(lang)) sbgrp_df <- sbgrp_df[sbgrp_df$lang == lang,]
#   sbgrp_df <- sbgrp_df[sbgrp_df$subj == subj,]
#   # for tests taken in english, model will include school and subgroup within school random effects,
#   # but for tests taken in spanish, almost all students are Hispanic/Latino, so
#   # the `all` and `ethh` scores are similar. Therefore we only include a school-level effect
#   re_form <- if (lang == "e") paste0("(1|", id_col, "/", sbgrp_col, ")") else paste0("(1|", id_col, ")")
#   form <- reformulate(c(Z_cols, "sg", re_form), response = W_col)
#   ctr <- do.call(glmmTMB::glmmTMBControl, list(...))
#   grades <- unique(sbgrp_df[[grade_col]][!is.na(sbgrp_df[[grade_col]])])
#   fits <- mapply(
#     function(j) {
#       fit_df <- sbgrp_df[sbgrp_df[[grade_col]] == j,]
#       init_sig_vals <- Reduce(
#         c,
#         vapply(
#           seq_len(nrow(fit_df)),
#           function(ix) get_csem(scr = fit_df$rs[ix], subj = subj
#                                 , lang = ifelse(subj == "m", "e", lang) # CSEM values are the same in english and spanish for math
#                                 , year = year, grade = j),
#           numeric(1L)
#         )
#       )
#       first_fit <- glmmTMB::glmmTMB(form, fit_df, weights = d, family = gaussian()
#                                     , dispformula = ~0-offset(log(init_sig_vals/sqrt(d)))
#                                     , control = ctr)
#       Xpred <- predict(first_fit, type = "response")
#       use_sig_vals <- Reduce(
#         c,
#         vapply(
#           Xpred,
#           function(rs) get_csem(scr = rs, subj = subj
#                                 , lang = ifelse(subj == "m", "e", lang), # CSEM values are the same in english and spanish for math
#                                 year = year, grade = j),
#           numeric(1L)
#         )
#       )
#       glmmTMB::glmmTMB(
#         form, fit_df[row.names(first_fit$frame),], weights = d, family = gaussian()
#         , dispformula = ~0-offset(log(use_sig_vals/sqrt(d))), control = ctr
#       )
#     },
#     grades,
#     SIMPLIFY = FALSE
#   )
#   
#   names(fits) <- paste0("g", grades)
#   return(fits)
# }

get_mle_logits <- function(logreg, eb_preds, scrs_yr, sgs, scr_colnms, K_mix_comps = 3) {
  # Fk_values comes from adsy_app_constants.R
  ski <- Fk_values[[K_mix_comps]][["ski"]]
  pki <- Fk_values[[K_mix_comps]][["pki"]]
  
  ## make a dataframe of empirical bayes predictions wide by tx_schoolid
  eb_preds <- widen_eb_preds(eb_preds, sgs, scrs_yr)
  
  # get test-taker counts (disaggregated by language for making a weighted avg.
  # of the eb preds) and other columns in the design matrix of the propensity score fit
  d_colnms <- c(
    gsub("_rs_", "_d_", gsub(paste0("_", scrs_yr), paste("", "e", scrs_yr, sep = "_"), scr_colnms)),
    gsub("_rs_", "_d_", gsub(paste0("_", scrs_yr), paste("", "s", scrs_yr, sep = "_"), scr_colnms)),
    gsub("_rs_", "_d_", scr_colnms)
  )
  dmat <- logreg$data[#expand.model.frame(logreg, c("tx_schoolid", d_colnms))[
    , c("tx_schoolid", d_colnms, "Ti",
        setdiff(gsub(", scale = FALSE)", "",
                     gsub("scale(", "", attr(terms(logreg), "term.labels"), fixed = TRUE),
                     fixed = TRUE), scr_colnms))
  ]
  dmat <- dplyr::left_join(dmat, eb_preds, by = "tx_schoolid")
  
  # get the test language-aggregated average scaled score
  wtd_scrs <- purrr::map_dfc(
    scr_colnms,
    wt.avg_eb_preds_by_lang,
    df = dmat,
    yr = scrs_yr
  )
  wtd_scrs <- dplyr::mutate_at(
    wtd_scrs,
    scr_colnms,
    ~ ifelse(.x == 0, mean(.x, na.rm = TRUE), .x)
  )
  dmat <- dplyr::bind_cols(dmat, wtd_scrs)

  # eb_preds <- eb_preds[, c("tx_schoolid", scr_cols)]
  # scr_colnms <- if (model_w_imputes <- all(paste(scr_base_colnms, "imputed", sep = "_") %in%
  #                                          names(logreg$coefficients))) {
  #   paste(scr_base_colnms, "imputed", sep = "_")
  # } else {
  #   scr_base_colnms
  # }
  # colnames(eb_preds) <- c("tx_schoolid", scr_colnms)

  # if (model_w_imputes) {
  #   dmat <- dmat[setdiff(seq_len(nrow(model.frame(logreg, na.action = na.pass))), na.action(logreg)),]
  #   dmat <- dplyr::bind_cols(
  #     dmat,
  #     purrr::map_dfc(
  #       scr_base_colnms,
  #       function(scr_prefix) {
  #         dplyr::summarize(
  #           dplyr::filter(dmat, !is.na(!!dplyr::sym(gsub("rs", "d", scr_prefix))) &
  #                           !!dplyr::sym(gsub("rs", "d", scr_prefix)) > 0),
  #           !!dplyr::sym(paste(scr_prefix, "grand_mean", sep = "_")) := mean(
  #             !!dplyr::sym(paste(scr_prefix, "imputed", sep = "_")), na.rm = TRUE)
  #         )
  #       }
  #     )
  #   )
  #   dmat[, paste(scr_base_colnms, "imputed", sep = "_")] <- purrr::map_dfc(
  #     scr_cols,
  #     function(scr_prefix) {
  #       dplyr::transmute_at(
  #         dmat,
  #         paste(scr_prefix, "imputed", sep = "_"),
  #         ~ ifelse(!is.na(!!dplyr::sym(gsub("rs", "d", scr_prefix))) &
  #                    !!dplyr::sym(gsub("rs", "d", scr_prefix)) > 0,
  #                  .x,
  #                  !!dplyr::sym(paste(scr_prefix, "grand_mean", sep = "_")))
  #       )
  #     }
  #   )
  # }
  
  sig_vals_by_lang <- mapply(
    function(scr_colnm, df, year) {
      lang <- ifelse(substr(scr_colnm, 1, 1) == "m", "e",
                     substr(scr_colnm,
                            regexpr(paste0("_", year), scr_colnm) - 1,
                            regexpr(paste0("_", year), scr_colnm) - 1))
      get_vec_csem(df[[scr_colnm]]
                   , rep(lang, nrow(df))
                   , substr(scr_colnm, 1, 1)
                   , substr(scr_colnm, regexpr("_g\\d{1}", scr_colnm) + 2, regexpr("_g\\d{1}", scr_colnm) + 2)
                   , year)
    },
    c(gsub(paste0("_", scrs_yr), paste("", "e", scrs_yr, sep = "_"), scr_colnms),
      gsub(paste0("_", scrs_yr), paste("", "s", scrs_yr, sep = "_"), scr_colnms)),
    MoreArgs = list(df = dmat, year = scrs_yr)
  )
  colnames(sig_vals_by_lang) <- gsub("_rs_", "_sig_", colnames(sig_vals_by_lang))

  dmat <- dplyr::bind_cols(
    dmat,
    purrr::map_dfc(
      gsub("_rs_", "_sig_", scr_colnms),
      wt.avg_meas_error_by_lang,
      df = cbind(dmat, sig_vals_by_lang),
      yr = scrs_yr
    )
  )
  sig_cols <- gsub("_rs_", "_sig_", scr_colnms)
  
  # use_sig_vals <- mapply(
  #   function(scr_prefix, dmat, year) {
  #     scr_col <- if (model_w_imputes) paste(scr_prefix, "imputed", sep = "_") else scr_prefix
  #     sig_vals <- vapply(dmat[,scr_col]
  #                        , get_csem
  #                        , numeric(1L)
  #                        , subj = substr(scr_prefix, 1, 1)
  #                        , lang = ifelse(substr(scr_prefix, 1, 1) == "m", "e",
  #                                        substr(scr_prefix,
  #                                               regexpr(paste0("_", year), scr_prefix) - 1,
  #                                               regexpr(paste0("_", year), scr_prefix) - 1))
  #                        , year = year
  #                        , grade = substr(scr_prefix, regexpr("_g\\d{1}", scr_prefix) + 2, regexpr("_g\\d{1}", scr_prefix) + 2))
  #     # if a structural zero, set sigma = 0
  #     sig_vals[is.na(dmat[[(gsub("rs", "d", scr_prefix))]]) |
  #                (!is.na(dmat[[(gsub("rs", "d", scr_prefix))]]) &
  #                   dmat[[(gsub("rs", "d", scr_prefix))]] == 0)] <- if (model_w_imputes) 0 else NA_real_
  #     sig_vals
  #   }
  #   , scr_cols
  #   , MoreArgs = list(dmat = dmat, year = scrs_yr)
  # )
  # mij <- dmat[, gsub("rs", "d", scr_cols)]
  # if (model_w_imputes) mij[use_sig_vals == 0] <- 1
  # varWij <- (use_sig_vals / sqrt(mij))^2

  coeffs_with_meas_err <- grepl(
    paste0("(", paste(scr_colnms, collapse = "|"), ")"),
    names(logreg$coefficients)
  )
  sig_cols_for_coeffs_with_meas_err <- vapply(
    names(logreg$coefficients)[coeffs_with_meas_err],
    function(nm) gsub("_rs_", "_sig_", names(which(sapply(scr_cols, function(col) grepl(col, nm))))),
    character(1L),
    USE.NAMES = FALSE
  )
  pihats <- rowSums(
    sweep(
      pnorm(
        sweep(
          replicate(K_mix_comps,
                    model.matrix(formula(logreg),
                                 model.frame(formula(logreg), dmat, na.action = na.pass)) %*% logreg$coefficients, simp = TRUE),
          2, ski, FUN = "*") /
          sqrt(1 + t(replicate(nrow(dmat), ski^2)) * rowSums(
            sweep(dmat[, sig_cols_for_coeffs_with_meas_err]^2,
                  2,
                  logreg$coefficients[coeffs_with_meas_err]^2,
                  FUN = "*")))
      ),
      2,
      pki,
      FUN = "*"
    )
  )
  
  return(log(pihats / (1-pihats)))
}

get_rc_logits <- function(lr_form, fit_df, eb_preds, scr_cols, sgs, scrs_yr) {
  ## make a dataframe of empirical bayes predictions wide by tx_schoolid
  eb_preds <- widen_eb_preds(eb_preds, sgs, scrs_yr)
  # eb_preds <- eb_preds[, c("tx_schoolid", scr_cols)]

  fit_df <- dplyr::left_join(
    fit_df[, c("tx_schoolid", setdiff(colnames(fit_df),
                                      c(colnames(eb_preds), scr_cols)))],
    eb_preds,
    by = "tx_schoolid"
  )
  wtd_scrs <- purrr::map_dfc(
    scr_cols,
    wt.avg_eb_preds_by_lang,
    df = fit_df,
    yr = scrs_yr
  )
  wtd_scrs <- dplyr::mutate_at(
    wtd_scrs,
    scr_cols,
    ~ ifelse(.x == 0, mean(.x, na.rm = TRUE), .x)
  )
  fit_df <- dplyr::bind_cols(fit_df, wtd_scrs)

  
  # fit_df <- dplyr::left_join(
  #   fit_df[, c("tx_schoolid", setdiff(colnames(fit_df), colnames(eb_preds)))],
  #   eb_preds,
  #   by = "tx_schoolid"
  # )
  rc_reg <- glm(lr_form, fit_df, family = binomial())

  # need the length to align with the original dataframe, so we insert the predictions
  # into a vector of the correct length
  # logits <- rep(NA_real_, nrow(fit_df))
  # logits[setdiff(seq_len(nrow(fit_df)), na.action(rc_reg))] <- predict(rc_reg, type = "link")

  # now we NA out the structural zeros (why would we keep a score that doesn't exist?)
  # logits[apply(fit_df[, gsub("rs", "d", scr_cols)],
  #              1,
  #              function(r) any(is.na(r) | (!is.na(r) & r == 0)))] <- NA 
  
  return(predict(rc_reg, type = "link"))
  # return(logits)
}

ps_fullmatch <- function(pmod, campus_df, fullmatch_params, ...) {
  mc <- match.call()
  if (!is.null(calip <- mc$caliper_width)) {
    # calip <- calip * sqrt(mean(tapply(predict(pmod, type = "link"), pmod$model$a, var)))
    calip <- calip
  }
  pdists <- suppressWarnings(optmatch::match_on(pmod, caliper = calip))
  pdists <- Reduce(`+`, ...)
  suppressWarnings(do.call(optmatch::fullmatch,
                           c(list(x = pdists, data = campus_df), fullmatch_params)))
}

fullmatch_logits <- function(df, trt_col, caliper_width, logits, fullmatch_params, ...) {
  a <- df[[trt_col]]
  if (inherits(logits, "numeric")) {
    dists <- optmatch::match_on(logits, z = a, caliper = caliper_width)
    # add the grades_match, which is passed into ...
    dists <- Reduce(`+`, c(list(dists, ...)))
    eval(
      do.call(call, c(list(name = "fullmatch", x = dists, data = df), fullmatch_params))
    )
  } else {
    warning("Taking logits to be a fitted propensity score model")
    eval(
      do.call(call, c(list(name = "ps_fullmatch", pmod = logits, campus_df = df,
                           caliper_width = caliper_width,
                           fullmatch_params = fullmatch_params),
                      list(...)))
    )
  }
}

# assess balance on Xhats and Zs
calc_balmetrics <- function(match, to_match_df, eb_preds, sgs, scrs_yr, scr_cols, trt_col, balcols) {
  # do the EB preds thing
  eb_preds <- widen_eb_preds(eb_preds, sgs, scrs_yr)
  to_match_df <- dplyr::left_join(
    to_match_df[, c("tx_schoolid", setdiff(colnames(to_match_df), c(colnames(eb_preds), scr_cols)))],
    eb_preds,
    by = "tx_schoolid"
  )
  wtd_scrs <- purrr::map_dfc(
    scr_cols,
    wt.avg_eb_preds_by_lang,
    df = to_match_df,
    yr = scrs_yr
  )
  wtd_scrs <- dplyr::mutate_at(
    wtd_scrs,
    scr_cols,
    ~ ifelse(.x == 0, mean(.x, na.rm = TRUE), .x)
  )
  colnames(wtd_scrs) <- scr_cols
  to_match_df <- dplyr::bind_cols(to_match_df, wtd_scrs)
  
  match_df <- dplyr::left_join(
    cbind(to_match_df, b = match), eb_preds, by = "tx_schoolid"
  )
  match_df$a <- as.numeric(match_df[[trt_col]])
  match_df <- match_df[!is.na(match_df$b),]

  ns <- rowsum(rep(1, nrow(match_df)), match_df$b)
  pib <- rowsum(match_df$a, match_df$b) / ns
  wts <- replicate(
    length(balcols),
    (match_df$a + (1-match_df$a) * pib[match_df$b,] / (1-pib[match_df$b,])) * ns[match_df$b,] / sum(ns)
  )

  wts[, balcols %in% scr_cols] <- wts[, balcols %in% scr_cols] * sapply(
    balcols[balcols %in% scr_cols],
    function(colnm, df) {
      if (grepl("_all_", colnm)) {
        gstr <- substr(colnm, regexpr("_g\\d{1}", colnm) + 1, regexpr("_g\\d{1}", colnm) + 2)
        g_bucket_cols <- grades_bucket_cols[grepl(gstr, grades_bucket_cols)]
        as.numeric(
          rowSums(df[, g_bucket_cols, drop=FALSE]) > 0 & if (length(unique(df$b)) > 1) {
            !(df$b %in% unique(df$b[df$a == 1 & rowSums(df[, g_bucket_cols, drop=FALSE]) == 0]))
          } else TRUE
        )
      } else {
        as.numeric(df[[gsub("_rs_", "_d_", colnm)]] > 0 & if (length(unique(df$b)) > 1) {
          !(df$b %in% unique(df$b[df$a == 1 & df[[gsub("_rs_", "_d_", colnm)]] == 0]))
        } else TRUE)
      }
    },
    df = match_df
  )
  trt_means <- colSums(match_df$a * wts * match_df[, balcols,drop=FALSE]) / colSums(match_df$a * wts)
  ctrl_means <- colSums((1-match_df$a) * wts * match_df[, balcols,drop=FALSE]) / colSums((1-match_df$a) * wts)
  sds <- sqrt(rowMeans(cbind(
    colSums(match_df$a * wts * match_df[, balcols,drop=FALSE]^2) / colSums(match_df$a * wts) -
      trt_means^2,
    colSums((1-match_df$a) * wts * match_df[, balcols,drop=FALSE]^2) / colSums((1-match_df$a) * wts) -
      ctrl_means^2
  )))
  
  c(setNames(abs(trt_means - ctrl_means), paste("matched", names(trt_means), sep = "_")),
    setNames(abs(trt_means - ctrl_means) / sds, paste("matched_stdized", names(trt_means), sep = "_")))
}

#' For a given matched sample, estimate (possibly multiple) ETT's and get
#' max-t adjusted p-values
#' columns
#' @param subj vector of subjects, one or both of `c("m", "r")`
#' @param grades list of grades, include `NULL` to estimate the marginal ETT
#'  across grades
#' @param analysis_yr `"19"` or `"21"`
#' @param analysis_sgs list of subgroups to estimate ETT's for, include `NULL`
#'  to estimate the marginal ETT across sgs
#' @param memod_fit_df dataframe that was used to fit mixed effects models. use
#'  this because it's already long with the desired regressors joined. we'll use
#'  it to make the fitting dataframes for the covariance adjustment and
#'  treatment effect regressions
#' @param schls_to_estimate treated schools to estimate ETT from 
#' @param ... Can include cmod_subset to specify the schools to use for fitting
#'  the covariance adjustment model
#' @return glht from multcomp package
estimate_trt_effect <- function(subj,
                                grades,
                                analysis_yr,
                                analysis_sgs,
                                covs,
                                schoolmatch,
                                memod_fit_df,
                                lr_data,
                                schls_to_estimate,
                                ...) {
  mc <- match.call()
  if (!inherits(analysis_sgs, "list")) analysis_sgs <- as.list(analysis_sgs)

  scrs_yrs <- if (analysis_yr == 19) 19 else c(21, 19)
  cmod_df <- Reduce(
    function(l, r) {
      # in the join function, we only keep test scores from previous years
      dplyr::left_join(l, r[,grepl("^(tx_schoolid|sg|m_d_|m_rs_|r_d_|r_rs_)", colnames(r))],
                       by = c("tx_schoolid", "sg"))
    },
    lapply(
      scrs_yrs,
      function(yr, df) {
        out <- tidyr::pivot_wider(df[df$join_yr == yr,]
                           , id_cols = setdiff(colnames(df),
                                               c("d", "rs", "subj", "grade", "lang", "join_yr"))
                           , names_from = c("subj", "grade", "lang", "join_yr")
                           , values_from = c("d", "rs")
                           , names_glue = paste("{subj}_{.value}_g{grade}_{lang}", yr, sep = "_"))
        agg_scrs <- purrr::map_dfc(
          paste(c("m_rs_g3", "r_rs_g3", "m_rs_g4", "r_rs_g4", "m_rs_g5", "r_rs_g5"), yr, sep = "_"),
          wt.avg_scrs_by_lang,
          df = out, yr = yr
        )
        agg_scrs <- dplyr::mutate_all(
          agg_scrs,
          ~ ifelse(.x == 0, NA_real_, .x)
        )
        out <- dplyr::bind_cols(out, agg_scrs)
        out
      },
      df = memod_fit_df
    )
  )
  cmod_df$uid <- cmod_df$tx_schoolid # this column needs to align with the Obs_Spec object that will be crated
  cmod_df$uniqid <- seq_len(nrow(cmod_df)) # unique id for the school/grade/subj/lang/sg combo to help propertee
  
  # store analysis_data now before we subset the cmod dataframe
  analysis_data <- cmod_df
  match_limits <- !is.na(schoolmatch) & schoolmatch %in% schoolmatch[
    lr_data$tx_schoolid %in% schls_to_estimate]

  # if a cmod_subset is provided, use it. otherwise, limit the covariance
  # adjustment model fit to the matched sets for treated schools specified in
  # schls_to_estimate
  if (!is.null(mc$cmod_subset)) {
    cmod_df <- subset(cmod_df, eval(mc$cmod_subset, cmod_df))
  } else {
    cmod_df <- subset(cmod_df, tx_schoolid %in% lr_data$tx_schoolid[match_limits])
  }

  # if assignment indicator not in cmod, drop treated schools from model fit
  if (!("a" %in% covs)) {
    cmod_df <- subset(cmod_df, a == 0)
  }
  cmod_df$a <- as.numeric(cmod_df$tx_schoolid %in% schls_to_estimate)

  if (!is.null(covs)) {
    cmods <- mapply(
      function(subj, grade, analysis_yr, cmod_df, covs) {
        rs_col <- paste(subj, "rs", paste0("g", grade), analysis_yr, sep = "_")
        wts_col <- paste(subj, "d", paste0("g", grade), analysis_yr, sep = "_")
        ca_form <- reformulate(covs, response = rs_col)
        cmod <- lm(ca_form, cmod_df, weights = cmod_df[[wts_col]])
        
        # this workaround allows lmitt() to find the correct information in the
        # temp environments created by mapply
        terms.env <- environment(terms(cmod))
        e <- new.env(parent = terms.env)
        e$ca_form <- ca_form
        e$fit_df <- cmod_df
        e$wts_col <- cmod_df[[wts_col]]
        environment(cmod$terms) <- e
        cmod$call$data <- quote(fit_df)
        cmod$call$weights <- quote(fit_df[[wts_col]])
        return(cmod)
      },
      rep(subj, length(grades)),
      rep(grades, each = length(subj)),
      MoreArgs = list(analysis_yr = analysis_yr,
                      cmod_df = cmod_df,
                      covs = covs),
      SIMPLIFY = FALSE
    )
    names(cmods) <- paste(rep(subj, length(grades)),
                          rep(grades, each = length(subj)),
                          sep = "_")
  } else {
    cmods <- NULL
  }

  rcols <- vapply(cmods,
                  function(x) deparse(formula(x)[[2L]], width.cutoff = 500L),
                  character(1L),
                  USE.NAMES = FALSE)
  tmods <- mapply(
    fit_lmitt,
    rep(rcols, each = length(analysis_sgs)),
    rep(analysis_sgs, length(cmods)),
    rep(cmods, each = length(analysis_sgs)),
    MoreArgs = list(match_df = lr_data
                    , match = schoolmatch
                    , treated_schools = schls_to_estimate
                    , analysis_df = analysis_data
                    , id_col = "tx_schoolid"
                    , trt_col = "Ti"
                    , uniqid_col = "uniqid"),
    SIMPLIFY = FALSE,
    USE.NAMES = FALSE
  )
  names(tmods) <- paste(rep(rcols, each = length(analysis_sgs)),
                        rep(analysis_sgs, length(cmods)), sep = "_")
  
  # make mmm object
  mmm.call <- do.call(
    call,
    c(list("mmm"),
      lapply(seq_along(tmods),
             function(ix) str2lang(paste0("quote(tmods[[", ix, "]])"))))
  )
  multobj <- eval(mmm.call)
  names(multobj) <- names(tmods)
  
  # make glht object
  linfct <- t(
    vapply(
      seq_along(tmods),
      function(ix) replace(numeric(length(tmods)*4L), 4L*(ix-1L)+2L, 1L),
      numeric(length(tmods)*4L)
    )
  )
  multcomp::glht(multobj, linfct, alternative = "two.sided", vcov. = vcov_tee,
                 type = "HC2", cov_adj_rcorrect = "HC1",
                 df = nrow(tmods[[1L]]$model) - 2L)
}

get_covs <- function(grade = NULL, pretrt_scrs_yr = NULL, include_trt = TRUE) {
  if (all(!is.null(c(grade, pretrt_scrs_yr)))) {
    covs <- paste(c("m", "r"), "rs", paste0("g", grade), pretrt_scrs_yr, sep = "_")
  } else if (!is.null(grade) | !is.null(pretrt_scrs_yr)) {
    stop("Must provide neither or both of grade and pretrt_scrs_yr")
  } else {
    covs <- numeric(0L)
  }
  covs <- c(covs,
            setdiff(ccd_keep_cols, "hispanic_total_enrollment_perc"),
            "sg",
            if (include_trt) "a" else NULL)
  return(covs)
}

fit_lmitt <- function(match_df, # for forming the StudySpecification object
                      match, # for forming the StudySpecification
                      treated_schools, # for forming the trt col in the analysis df
                      analysis_df, # for forming the df passed to lmitt()
                      id_col, # school ID column that exists in both match_df and scrs_df
                      trt_col, # treatment column in match_df; simplifies forming the StudySpecification
                      response_col, # response column in scrs_df
                      sg = NULL, # subgroup to estimate ETT for; if `NULL`, marginal ETT
                      cmod = NULL, # covariance adjustment model, if desired
                      uniqid_col = NULL, # need if cmod provided because may have multiple obs from the same school (across years and subgroups, for example)
                      limit_match = NULL) { # can limit matches if need be
  limit_match <- limit_match %||% TRUE
  # des_df <- data.frame(
  #   uid = match_df[[id_col]],
  #   blk = match,
  #   a = as.numeric(match_df[[trt_col]])
  # )[!is.na(match),]
  des_df <- data.frame(
    uid = match_df[[id_col]],
    blk = match,
    a = as.numeric(match_df[[trt_col]])
  )[!is.na(match) & limit_match,]
  des <- propertee::obs_spec(a ~ unitid(uid) + block(blk), des_df)
  
  if (!is.null(cmod)) {
    expand_cols <- c(id_col, "sg", uniqid_col)
    analysis_df <- dplyr::left_join(
      analysis_df[, setdiff(colnames(analysis_df), uniqid_col)],
      expand.model.frame(cmod, expand_cols)[, expand_cols, drop=FALSE],
      by = setdiff(expand_cols, uniqid_col)
    )
    need_cols <- setdiff(attr(terms(cmod), "term.labels"), colnames(analysis_df))
    if (length(need_cols) > 0) {
      analysis_df <- dplyr::left_join(
        analysis_df, expand.model.frame(cmod, uniqid_col)[,c(uniqid_col, need_cols)], by = uniqid_col
      )
    }
  }

  analysis_df$uid <- analysis_df[[id_col]]
  analysis_df$a <- as.numeric(analysis_df$uid %in% treated_schools)
  analysis_df <- analysis_df[analysis_df$uid %in% des_df$uid,] # this ensures we only keep ID's from limit_match
  
  form <- reformulate("1", response = response_col)
  if (!is.null(cmod)) {
    os <- cov_adj(cmod, analysis_df, des, by = uniqid_col)
  } else {
    os <- NULL
  }
  wts_col <- gsub("_rs_", "_d_", response_col)
  sg_ix <- if (is.null(sg)) 1 else analysis_df$sg == sg
  wts <- ett(des, data = analysis_df) * analysis_df[[wts_col]] * as.numeric(sg_ix)
  fit <- propertee::lmitt(form, analysis_df, spec = des, absorb = FALSE,
                          weights = wts, offset = os)
  return(fit)
}

if (sys.nframe() == 0) {
  source("adsy_app_constants.R")
  data_dir <- "adsy_data"
  load(file.path(data_dir, "ccd.Rdata"))
  load(file.path(data_dir, "adsy_schools_20_21.Rdata"))
  
  scrs_yr <- 19 # data collected at the end of the academic year
  ccd_yr <- 18 # data collected at the beginning of the academic year
  grades <- seq(3, 5)
  sgs_keep <- c("all", "ethh", "ethb", "etha", "ethw")
  join_cols <- c("tx_schoolid", "year")
  get_scr_cols <- function(subj, sg) paste(subj, sg, c("d", "rs"), sep = "_")
  ccd_keep_cols <- c("sch_type", "charter", "title1", "magnet",
                     paste(c("white", "hispanic", "african_american", "asian", "frl", "female"),
                           "total_enrollment_perc",
                           sep = "_"))
  
  # form long dataframe of subgroup scores that can be used to fit a model for empirical bayes predictions
  scrs <- Reduce(
    rbind,
    lapply(
      grades,
      function(g, dir, sgs, join_cols) {
        e <- new.env(parent = emptyenv())
        load(file.path(dir, paste0("g", g, "_scores.Rdata")), envir = e)
        obj_nm <- paste0("g", g)
        e[[obj_nm]][["grade"]] <- g
        Reduce(
          rbind,
          lapply(sgs, function(sg) {
            rbind(setNames(cbind(e[[obj_nm]][, c(join_cols, "lang", "grade", get_scr_cols("m", sg))],
                                 "m", sg),
                           c(join_cols, "lang", "grade", c("d", "rs"), "subj", "sg")),
                  setNames(cbind(e[[obj_nm]][, c(join_cols, "lang", "grade", get_scr_cols("r", sg))],
                                 "r", sg),
                           c(join_cols, "lang", "grade", c("d", "rs"), "subj", "sg")))
          })
        )
      },
      dir = data_dir, sgs = sgs_keep, join_cols = join_cols
    )
  )

  wide_scrs <- tidyr::pivot_wider(scrs[scrs$year == scrs_yr,], id_cols = "tx_schoolid"
                                  , values_from = c("d", "rs")
                                  , names_from = c("year", "lang", "grade", "subj", "sg")
                                  , names_glue = "{subj}_{sg}_{.value}_g{grade}_{lang}_{year}")
  ccd_wide <- tidyr::pivot_wider(ccd[ccd$year == ccd_yr,]
                                 , id_cols = "tx_schoolid"
                                 , values_from = setdiff(colnames(ccd), join_cols)
                                 , names_from = "year")
  
  ## make treatment column (note 5/19/25: only using schools that added days
  ## before spring 2021 testing)
  treated_to_assess <- SCHLS_FOR_ESTIMATES
  all_treated <- vapply(
    c(adsy_20_21_raw$campus, 31901101, 31901104, 31901108, 31901117, 31901121,
      31901122, 31901123, 31901129, 31901132, 31901137, 31901140, 31901143,
      57803116, 71907106, 94902109, 177901001),
    function(id) paste(c(rep("0", 9-nchar(id)), id), collapse = ""),
    character(1L)
  )
  treated_to_exclude <- setdiff(all_treated, treated_to_assess)
  treated_missing_scrs <- setdiff(treated_to_assess, unique(scrs$tx_schoolid[scrs$tx_schoolid %in% treated_to_assess & scrs$year == scrs_yr]))
  treated_to_assess <- replace(
    replace(treated_to_assess, treated_to_assess == "071908101", "071908044"), # this should be the id for TORNILLO INT
    treated_to_assess == "015825101", "015825001") # change LIGHTHOUSE EL to LIGHTHOUSE MIDDLE because its test scores are stored under LIGHTHOUSE MIDDLE
  treated_to_exclude <- replace(
    replace(treated_to_exclude, treated_to_exclude == "071908101", "071908044"), # this should be the id for TORNILLO INT
    treated_to_exclude == "015825101", "015825001") # change LIGHTHOUSE EL to LIGHTHOUSE MIDDLE because its test scores are stored under LIGHTHOUSE MIDDLE
  
  # from NCES.gov we can confirm the remaining treated schools don't serve 3-5 graders, or we can check scrs to see if
  # there are no rows where year==scrs_yr: 068901117 = PEASE ELEMENTARY (K-2), 068901106 = CARVER EARLY Ee-kg (PK),
  # 068901115 = LAMAR EARLY Ee-kg (PK), 068901123 = EL MAGNET AT ZAVALA (K-2), 015908115 has no testing pre-2021,
  # 057813105 has no testing pre-2021, 128901107 = KARNES CITY PRIMARY (PK-1), 058906001 has no testing pre-2021
  # 082903104 = TED FLORES ELEMENTARY (PK-2), 205904101 = MATHIS ELEMENTARY (PK-2), 125903101 = ORANGE GROVE PRIMARY (PK-2)
  # 205905103 = ODEM ELEMENTARY (PK-2), 057851101 has no testing pre-2021, 108807138 and 108807139 have no testing pre-2021,
  # 123910133 = LUCAS PK (serves PK only; found this in an Excel sheet from a google search of this school ID),
  # 123910132 = BINGMAN PK (PK), 177901102 = ROSCOE COLLEGIATE MONTESSORI EARLY CHILDHOOD (PK-K),
  # 057803018 and 057803017 have no pre-2021 testing, 071801004 has no pre-2021 testing,
  # 126907041 = RIO VISTA MIDDLE (6-8), 057909181 = PARSONS PRE-K CENTER (PK), 057909180 = CISNEROS PRE-K CENTER (PK),
  # 152806001 has no pre-2021 testing, 015807004 = SOUTHWEST PREP SCHOOL NORTHWEST (6-12),
  # 057803116 has no pre-2021 testing, 177901001 = ROSCOE COLLEGIATE HS (6-12)
  
  # make the dataframe for fitting the logistic regression
  lr_fit_df <- dplyr::left_join(
    wide_scrs,
    ccd_wide[, c("tx_schoolid", paste(c(ccd_keep_cols, "sy_status"), ccd_yr, sep = "_"))],
    by = "tx_schoolid"
  )
  lr_fit_df$Ti <- as.numeric(lr_fit_df$tx_schoolid %in% treated_to_assess)
  
  ## drop the treated to schools to exclude.
  ## additionally, most schools in the KIPP TEXAS PUBLIC SCHOOLS ISD (district 227820) have no
  ## CCD info for the 18-19 school year and are marked by sy_status == 5. We exclude
  ## them from the ccd df because these missing values are coded as 0's rather than NA's,
  ## so these rows would not be dropped from the model fit. No other schools in this
  ## year have sy_status == 5. sy_status == 2 also indicates a closed school, so we remove
  ## the one control school with sy_status == 2. We also remove 2 control schools in the
  ## scores data that do not appear in the CCD data and 7 control schools that have no
  ## scores data. Finally, we remove, 5 control schools that do not have a measure of
  ## campus-wide free or reduced-price lunch eligibility
  lr_fit_df <- lr_fit_df[
    !(lr_fit_df[[paste("sy_status", ccd_yr, sep = "_")]] %in% c(2, 5)) &
      !(lr_fit_df$tx_schoolid %in% c(treated_to_exclude, "152504101", "221901124",
                                     "247903002", "221901008", "216901180",
                                     "041902180", "161914127", "101914051",
                                     "015907020", "057905099", "068901005",
                                     "071902173", "101906042", "165901020")),]
  
  # create a weighted average of scores on english and spanish language tests
  scr_cols <- apply(expand.grid(c("m", "r"), sgs_keep, paste0("rs_g", grades), scrs_yr),
                    1, function(...) paste(..., collapse = "_"))
  lr_fit_df <- dplyr::bind_cols(
    lr_fit_df,
    purrr::map_dfc(
      scr_cols,
      wt.avg_scrs_by_lang,
      yr = scrs_yr, df = lr_fit_df
    )
  )
  
  # from working with the student data, we think that school 227806037 is miscounted
  # in the aggregate data. there are no students in the student data taking 3rd-5th
  # grade tests, so we manually set those counts to 0 in the aggregated data
  lr_fit_df[lr_fit_df$tx_schoolid == "227806037", c(scr_cols, gsub("_rs_", "_d_", scr_cols))] <- 0
  
  # grand mean impute scores for schools missing English and Spanish language scores
  # (in wt.avg_scrs_by_lang, we'd coded them as 0)
  lr_fit_df <- dplyr::mutate_at(
    lr_fit_df,
    scr_cols,
    ~ ifelse(.x == 0, mean(.x, na.rm = TRUE), .x)
  )
  
  # # if we're going with the simpler model, drop the spanish language columns
  # scr_cols <- scr_cols[grepl("e_19$", scr_cols)]
  # # if using "ethb", don't use scores on spanish language test because no schools have them
  # scr_cols <- setdiff(scr_cols,
  #                     apply(expand.grid(c("m", "r"), "ethb", paste0("rs_g", grades), "s", scrs_yr),
  #                           1, function(...) paste(..., collapse = "_")))
  # lr_fit_df <- dplyr::mutate_at(
  #   lr_fit_df, scr_cols, list(missing = ~ ifelse(is.na(.x), 1, 0),
  #                             grand_mean = function(x) mean(x, na.rm = TRUE))
  # )
  # lr_fit_df <- lr_fit_df[
  #   rowSums(lr_fit_df[, paste(scr_cols, "missing", sep = "_")]) <=
  #     length(scr_cols) - 2,]
  # lr_fit_df <- dplyr::bind_cols(
  #   lr_fit_df,
  #   purrr::map_dfc(
  #     scr_cols,
  #     function(x) dplyr::transmute(lr_fit_df,
  #                                  !!dplyr::sym(paste(x, "imputed", sep = "_")) :=
  #                                    ifelse(!!dplyr::sym(paste(x, "missing", sep = "_")) == 1,
  #                                           !!dplyr::sym(paste(x, "grand_mean", sep = "_")),
  #                                           !!dplyr::sym(x)))
  #   )
  # )
  # english language math and reading missingness within grade are perfectly correlated, so
  # use just one indicator to indicate whether there are students in the grade
  # lr_fit_df <- dplyr::bind_cols(
  #   lr_fit_df,
  #   purrr::map2_dfc(
  #     rep(sgs, length(grades)),
  #     rep(grades, each = length(sgs)),
  #     function(sg, g, year) {
  #       missing_col <- paste(sg, "rs", paste0("g", g), "e", year, "missing", sep = "_")
  #       dplyr::transmute(
  #         lr_fit_df,
  #         !!dplyr::sym(missing_col) :=
  #           as.numeric(!!dplyr::sym(paste("m", missing_col, sep = "_")) + !!dplyr::sym(paste("r", missing_col, sep = "_")) >= 1)
  #       )
  #     },
  #     year = scrs_yr
  #   )
  # )
  
  # interact scores with grades served terms
  lr_fit_df <- dplyr::bind_cols(
    lr_fit_df,
    purrr::map_dfc(
      grades,
      function(g, df, yr) dplyr::transmute(
        df,
        !!dplyr::sym(paste0("serves_g", g)) :=
          replace_na(!!dplyr::sym(paste("m", "all", "d", paste0("g", g), yr, sep = "_")), 0) + 
            replace_na(!!dplyr::sym(paste("r", "all", "d", paste0("g", g), yr, sep = "_")), 0) > 0
      ),
      df = lr_fit_df,
      yr = scrs_yr
    )
  )
  grade_combs <- Reduce(
    cbind,
    mapply(
      function(m) {
        combs <- combn(grades, m)
        rbind(combs, matrix(0, nrow = length(grades) - m, ncol = ncol(combs)))
      },
      seq_along(grades)
    )
  )
  grades_bucket_cols <- apply(
    grade_combs,
    2,
    function(c) paste("is", paste(paste0("g", c[c != 0]), collapse = "_"), "schl", sep = "_")
  )
  lr_fit_df <- dplyr::bind_cols(
    lr_fit_df,
    purrr::map_dfc(
      grades_bucket_cols,
      function(col, df) {
        gs <- strsplit(gsub("_schl$", "", gsub("^is_", "", col)), "_")[[1L]]
        dplyr::transmute(
          df,
          !!dplyr::sym(col) :=
            as.numeric(
              dplyr::if_all(.cols = paste("serves", gs, sep = "_")
                            , .fns = ~ .x) & 
                if (length(gs) == length(grades)) TRUE else {
                  dplyr::if_all(.cols = paste("serves", setdiff(paste0("g", grades), gs), sep = "_")
                                , .fns = ~ !.x)
                }
            )
        )
      },
      df = lr_fit_df
    )
  )
  
  # last thing is to remove schools with is_g4_schl = 1 from the fit; there are only
  # 6 of them, all controls, 5 of which already have NA scores. this way, we don't
  # have to fit interaction terms for is_g4_schl, knowing they can't be matched to
  # treated schools anyways (since they don't serve all the grades treated schools do)
  lr_fit_df <- lr_fit_df[lr_fit_df$is_g4_schl == 0,]
  
  # count schools in each bucket
  bucket_cts <- colSums(lr_fit_df[, grades_bucket_cols])
  cat(paste(paste(names(bucket_cts), bucket_cts, sep = ": "), collapse = "\n\n"))

  # save the dataframe for fitting the propensity score model so we can use it on
  # cochran to fit the model that imputes masked values with student data
  saveRDS(lr_fit_df, file = "lr_fit_df.Rdata")
  
  # there are discrepancies between subgroup counts in the aggregate data and the
  # student-level data. we treat the student-level data as truth, so for schools
  # with 0 students in the student-level data belonging to a race/ethnicity group,
  # we set their student count to 0 in the aggregate data and additionally set
  # their scores to be the appropriate grand mean imputation. the specific ID'
  # are printed in the output of the student fit script. after making those changes,
  # we fit the PS model here, then fit the PS model using the student data with the
  # terms kept in the model selection process here.
  lr_fit_df[c(8, 24, 26, 205, 275, 282, 353, 493, 502, 512, 514, 519, 546, 610,
              630, 713, 755, 823, 836, 932, 1001, 1016, 1138, 1303, 1449, 1512,
              1520, 1577, 1601, 1806, 1932, 1990, 2087, 2106, 2234, 2249, 2272,
              2484, 2490, 2584, 2762, 2820, 2833, 2931, 2985, 3024, 3040,
              3074, 3103, 3117, 3198, 3237, 3256, 3260, 3343, 3377, 3383, 3390,
              3397, 3405, 3639, 3838, 3927, 3955, 3975, 4007, 4052, 4074, 4088,
              4127, 4181, 4330, 4338, 4346),
            c(gsub("_rs_", "_d_", scr_cols[grepl("ethb", scr_cols) & grepl("g3", scr_cols)]),
              scr_cols[grepl("ethb", scr_cols) & grepl("g3", scr_cols)])] <- dplyr::bind_cols(
                data.frame(m_ethb_d_g3_19 = 0, r_ethb_d_g3_19 = 0),
                dplyr::distinct(lr_fit_df[lr_fit_df$m_ethb_d_g3_19 == 0, "m_ethb_rs_g3_19"]),
                dplyr::distinct(lr_fit_df[lr_fit_df$r_ethb_d_g3_19 == 0, "r_ethb_rs_g3_19"])
            )
  lr_fit_df[c(845, 1627, 1744, 1920, 2087, 2912, 3478, 3517, 3628, 3923, 3925, 4031, 4362),
            c(gsub("_rs_", "_d_", scr_cols[grepl("ethh", scr_cols) & grepl("g3", scr_cols)]),
              scr_cols[grepl("ethh", scr_cols) & grepl("g3", scr_cols)])] <- dplyr::bind_cols(
              data.frame(m_ethh_d_g3_19 = 0, r_ethh_d_g3_19 = 0),
              dplyr::distinct(lr_fit_df[lr_fit_df$m_ethh_d_g3_19 == 0, "m_ethh_rs_g3_19"]),
              dplyr::distinct(lr_fit_df[lr_fit_df$r_ethh_d_g3_19 == 0, "r_ethh_rs_g3_19"])
            )
  lr_fit_df[c(55, 122, 173, 193, 293, 374, 407, 559, 615, 828, 1019, 1073, 1091,
              1134, 1206, 1492, 2087, 2514, 2557, 2596, 2642, 2697, 2822, 2880,
              2978, 2984, 3074, 3091, 3095, 3236, 3257, 3266, 3334, 3370, 3396,
              3445, 3673, 3848, 3956, 3968, 3976, 3992, 4052, 4074, 4091, 4113,
              4116, 4375, 4407, 4432, 4447),
            c("m_ethb_d_g4_19", "m_ethb_rs_g4_19")] <- dplyr::bind_cols(
              data.frame(m_ethb_d_g4_19 = 0),
              dplyr::distinct(lr_fit_df[lr_fit_df$m_ethb_d_g4_19 == 0, "m_ethb_rs_g4_19"])
            )
  lr_fit_df[c(55, 122, 173, 193, 293, 374, 407, 559, 615, 828, 1019, 1073, 1091,
              1134, 1206, 1492, 2087, 2514, 2557, 2596, 2642, 2697, 2822, 2880,
              2978, 2984, 3074, 3091, 3095, 3236, 3257, 3266, 3334, 3370, 3396,
              3445, 3673, 3956, 3968, 3976, 3992, 4052, 4074, 4091, 4113, 4116,
              4375, 4407, 4432, 4447),
            c("r_ethb_d_g4_19", "r_ethb_rs_g4_19")] <- dplyr::bind_cols(
              data.frame(r_ethb_d_g4_19 = 0),
              dplyr::distinct(lr_fit_df[lr_fit_df$r_ethb_d_g4_19 == 0, "r_ethb_rs_g4_19"])
            )
  lr_fit_df[c(2428, 2766, 3408, 3922, 3923, 4031, 4051, 4093, 4221, 4385, 4409, 4432),
            c(gsub("_rs_", "_d_", scr_cols[grepl("ethh", scr_cols) & grepl("g4", scr_cols)]),
              scr_cols[grepl("ethh", scr_cols) & grepl("g4", scr_cols)])] <- dplyr::bind_cols(
                data.frame(m_ethh_d_g4_19 = 0, r_ethh_d_g4_19 = 0),
                dplyr::distinct(lr_fit_df[lr_fit_df$m_ethh_d_g4_19 == 0, "m_ethh_rs_g4_19"]),
                dplyr::distinct(lr_fit_df[lr_fit_df$r_ethh_d_g4_19 == 0, "r_ethh_rs_g4_19"])
            )
  lr_fit_df[c(2, 20, 24, 27, 30, 50, 52, 121, 131, 144, 161, 163, 164, 166, 175,
              182, 186, 193, 200, 201, 224, 225, 230, 237, 238, 244, 262, 267,
              287, 295, 328, 381, 382, 391, 407, 491, 495, 499, 501, 505, 507,
              541, 572, 636, 702, 741, 819, 823, 825, 867, 988, 994, 1006, 1007,
              1009, 1027, 1029, 1087, 1095, 1184, 1257, 1278, 1306, 1318, 1322,
              1349, 1399, 1404, 1441, 1507, 1527, 1543, 1556, 1562, 1564, 1575,
              1640, 1641, 1644, 1645, 1646, 1795, 1796, 1838, 1842, 1895, 1942,
              1946, 1988, 1990, 2058, 2071, 2096, 2138, 2143, 2190, 2209, 2428,
              2438, 2449, 2487, 2501, 2596, 2614, 2623, 2625, 2643, 2659, 2697,
              2700, 2709, 2762, 2804, 2816, 2820, 2835, 2851, 2859, 2873, 2915,
              3052, 3054, 3067, 3073, 3077, 3085, 3094, 3103, 3122, 3192, 3199,
              3233, 3240, 3259, 3276, 3302, 3326, 3343, 3348, 3358, 3373, 3377,
              3382, 3383, 3406, 3448, 3457, 3458, 3463, 3469, 3475, 3534, 3545,
              3547, 3601, 3609, 3639, 3654, 3665, 3670, 3705, 3805, 3882, 3883,
              3889, 3890, 3892, 3901, 3913, 3923, 3924, 3925, 3927, 3946, 3958,
              3959, 3972, 3975, 3987, 3989, 3998, 4027, 4052, 4091, 4093, 4097,
              4108, 4109, 4113, 4137, 4158, 4162, 4188, 4192, 4207, 4218, 4225,
              4245, 4253, 4262, 4277, 4307, 4338, 4346, 4355, 4356, 4357, 4377,
              4381, 4400, 4406, 4408, 4423, 4432, 4446, 4461, 4464, 4465, 4483,
              4484, 4485, 4488, 4506, 4514, 4518, 4538, 4578, 4584, 4625, 4634,
              4640, 4669, 4687, 4688),
            c("m_ethb_d_g5_19", "m_ethb_rs_g5_19")] <- dplyr::bind_cols(
              data.frame(m_ethb_d_g5_19 = 0),
              dplyr::distinct(lr_fit_df[lr_fit_df$m_ethb_d_g5_19 == 0, "m_ethb_rs_g5_19"])
            )
  lr_fit_df[c(2, 20, 24, 27, 30, 50, 52, 121, 131, 144, 161, 163, 164, 166, 175,
              182, 186, 193, 200, 201, 224, 225, 230, 236, 237, 238, 244, 262,
              267, 287, 295, 328, 382, 391, 407, 410, 491, 495, 499, 501, 505,
              507, 541, 572, 636, 702, 741, 819, 823, 825, 831, 867, 988, 994,
              1006, 1007, 1009, 1027, 1029, 1087, 1095, 1184, 1194, 1257, 1278,
              1318, 1322, 1349, 1399, 1441, 1507, 1527, 1543, 1556, 1562, 1564,
              1575, 1640, 1641, 1644, 1645, 1646, 1795, 1796, 1893, 1895, 1942,
              1946, 1988, 1990, 2058, 2071, 2087, 2096, 2125, 2138, 2143, 2190,
              2209, 2449, 2487, 2501, 2596, 2614, 2623, 2625, 2643, 2659, 2697,
              2700, 2709, 2762, 2801, 2804, 2816, 2820, 2835, 2851, 2859, 2873, 
              2891, 2915, 3052, 3054, 3067, 3073, 3077, 3085, 3094, 3103, 3122,
              3192, 3199, 3233, 3240, 3259, 3276, 3302, 3326, 3343, 3348, 3358,
              3373, 3377, 3382, 3383, 3406, 3448, 3457, 3458, 3463, 3469, 3475,
              3545, 3547, 3601, 3609, 3639, 3654, 3665, 3670, 3705, 3805, 3882,
              3883, 3889, 3890, 3892, 3901, 3913, 3923, 3924, 3925, 3927, 3946,
              3958, 3972, 3975, 3987, 3989, 3998, 4027, 4052, 4091, 4093, 4097,
              4108, 4109, 4113, 4137, 4158, 4162, 4188, 4192, 4207, 4218, 4225,
              4228, 4245, 4253, 4262, 4277, 4307, 4338, 4343, 4346, 4355, 4356,
              4377, 4381, 4400, 4406, 4423, 4432, 4446, 4464, 4465, 4483, 4484,
              4485, 4488, 4506, 4514, 4518, 4538, 4578, 4584, 4625, 4634, 4640,
              4669, 4687, 4688),
            c("r_ethb_d_g5_19", "r_ethb_rs_g5_19")] <- dplyr::bind_cols(
              data.frame(r_ethb_d_g5_19 = 0),
              dplyr::distinct(lr_fit_df[lr_fit_df$r_ethb_d_g5_19 == 0, "r_ethb_rs_g5_19"])
            )
  lr_fit_df[c(62, 223, 631, 810, 893, 1618, 1631, 1798, 1876, 2488, 2489, 2830,
              2915, 3123, 3130, 3301, 3310, 3348, 3466, 3842, 3853, 3922, 3924,
              3927, 4052, 4086, 4087, 4097, 4377, 4385, 4407, 4486, 4489, 4492,
              4507, 4519, 4543, 4588, 4598, 4689, 4690),
            c("m_ethh_d_g5_19", "m_ethh_rs_g5_19")] <- dplyr::bind_cols(
              data.frame(m_ethh_d_g5_19 = 0),
              dplyr::distinct(lr_fit_df[lr_fit_df$m_ethh_d_g5_19 == 0, "m_ethh_rs_g5_19"])
            )
  lr_fit_df[c(62, 223, 631, 632, 810, 893, 1618, 1624, 1631, 1798, 1876, 2488,
              2489, 2526, 2765, 2830, 2859, 2891, 2915, 3123, 3130, 3301, 3310,
              3348, 3466, 3842, 3853, 3924, 3927, 4052, 4086, 4087, 4097, 4385,
              4407, 4464, 4489, 4492, 4507, 4519, 4543, 4588, 4598, 4689, 4690),
            c("r_ethh_d_g5_19", "r_ethh_rs_g5_19")] <- dplyr::bind_cols(
              data.frame(r_ethh_d_g5_19 = 0),
              dplyr::distinct(lr_fit_df[lr_fit_df$r_ethh_d_g5_19 == 0, "r_ethh_rs_g5_19"])
            )
  
  # the propensity score model will exclude interactions with is_g3_g4_g5_schl
  # to make it the reference level; we also don't include is_g3_g5_schl and is_g4_schl
  # because those schools have NA scores
  # lr_form <- reformulate(
  #   c(paste(setdiff(ccd_keep_cols, "sch_type"), ccd_yr, sep = "_"), # get rid of sch_type for PS model bc 121/125 treated_to_assess schools are standard ed
  #     # paste0("scale(", ccd_keep_cols[grepl("_perc$",ccd_keep_cols )], "_", ccd_yr, ", scale=FALSE)")
  #     paste(rep(scr_cols[grepl("eth(b|h)", scr_cols)],
  #               length(setdiff(grades_bucket_cols, c("is_g4_schl", "is_g3_g5_schl", "is_g3_g4_g5_schl")))),
  #           rep(setdiff(grades_bucket_cols,  c("is_g4_schl", "is_g3_g5_schl", "is_g3_g4_g5_schl")),
  #               each = sum(grepl("eth(b|h)", scr_cols))), sep = "*")),
  #   response = "Ti"
  # )
  # colSums(is.na(model.matrix(lr_form, model.frame(lr_form, lr_fit_df, na.action = na.pass))))

  # fit logistic regression
  lr_form <- reformulate(
    c(paste(c("charter", "magnet", "title1"), ccd_yr, sep = "_"),
      paste0("scale(",
             paste(setdiff(ccd_keep_cols[grepl("_perc$",ccd_keep_cols)],
                           "hispanic_total_enrollment_perc"), ccd_yr, sep = "_"),
             ", scale=FALSE)"),
      # paste(paste0("scale(", rep(scr_cols[grepl("eth(b|h)", scr_cols)], each = 3), ", scale = FALSE)"),
      #       rep(c("is_g3_schl", "is_g4_g5_schl"),
      #           sum(grepl("eth(b|h)", scr_cols))), sep = "*"),
      # "is_g3_g4_schl"),
      paste0("scale(", scr_cols[grepl("eth(b|h)", scr_cols)], ", scale = FALSE)")),
    response = "Ti"
  )
  saveRDS(attr(terms(lr_form), "term.labels"), file = "lr_fit_vars.Rdata")
  lr <- glm(lr_form, lr_fit_df, family = binomial())
  
  # here we see if we can reduce the model at all by testing whether interaction
  # terms significantly increase the likelihood
  # terms_to_test <- c(
  #   paste0("m_all_rs_g5_", scrs_yr, ":is_g4_g5_schl"), paste0("r_all_rs_g5_", scrs_yr, ":is_g4_g5_schl"),
  #   paste0("m_all_rs_g4_", scrs_yr, ":is_g4_g5_schl"), paste0("r_all_rs_g4_", scrs_yr, ":is_g4_g5_schl"),
  #   paste0("m_all_rs_g4_", scrs_yr, ":is_g3_g4_schl"), paste0("r_all_rs_g4_", scrs_yr, ":is_g3_g4_schl"),
  #   paste0("m_all_rs_g3_", scrs_yr, ":is_g3_g4_schl"), paste0("r_all_rs_g3_", scrs_yr, ":is_g3_g4_schl"),
  #   paste0("m_all_rs_g5_", scrs_yr, ":is_g5_schl"), paste0("r_all_rs_g5_", scrs_yr, ":is_g5_schl"),
  #   paste0("m_all_rs_g3_", scrs_yr, ":is_g3_schl"), paste0("is_g3_schl:r_all_rs_g3_", scrs_yr)
  # )
  # lr_fit_lrts <- drop1(lr, terms_to_test, test = "LRT")
  # lr_fit_lrts
  # 
  # 
  # # keep interactions that are significant at level 0.1
  # lr_form <- reformulate(
  #   c(paste(ccd_keep_cols, ccd_yr, sep = "_"), scr_cols,
  #     setdiff(grades_bucket_cols, c("is_g4_schl", "is_g3_g5_schl", "is_g3_g4_g5_schl")),
  #     row.names(lr_fit_lrts[lr_fit_lrts$`Pr(>Chi)`<=0.1 & row.names(lr_fit_lrts) != "<none>",])),
  #   response = "Ti"
  # )
  # lr <- glm(lr_form, lr_fit_df, family = binomial())
  
  # first restrict matches to controls that serve all the grades an ADSY school does.
  # control that serves grade 3 can match to a treatment that does not serve grade 3,
  # but a control that doesn't serve grade 3 can't match to a treatment that does serve grade 3
  grades_match <- Reduce(
    `+`,
    mapply(
      function(colnm, df) {
        gstr <- strsplit(gsub("_schl", "", gsub("^is_", "", colnm)), "_")[[1L]]
        # set all values of the vector to match on to TRUE, and just replace control
        # schools that don't serve all the grades a treated school does to FALSE.
        # this is because we don't care if control schools don't serve a grade that
        # a treated school likewise doesn't serve
        match_col <- rep(TRUE, nrow(df))
        match_col[rowSums(
          df[, grades_bucket_cols[rowSums(sapply(gstr, grepl, grades_bucket_cols)) ==
                                    length(gstr)], drop=FALSE]
        ) == 0 & !df$Ti] <- FALSE
        optmatch::exactMatch(setNames(match_col, row.names(df)), df$Ti)
      },
      grades_bucket_cols,
      MoreArgs = list(df = as.data.frame(lr_fit_df)),
      SIMPLIFY = FALSE
    )
  )
  
  # also restrict matches to being of the same school type
  sch_type_match <- optmatch::exactMatch(setNames(lr_fit_df$sch_type_18, row.names(lr_fit_df)),
                                         lr_fit_df$Ti)

  # fit models for math and reading in grades 3-5, first with all schools then
  # with only schools that have students taking the test in both English and Spanish
  memod_fit_df <- dplyr::left_join(
    dplyr::mutate(scrs, join_yr = as.numeric(year)),
    dplyr::mutate(ccd[,c(join_cols, ccd_keep_cols)], join_yr = as.numeric(year) + 1),
    by = c("tx_schoolid", "join_yr")
  )
  mods_list <- mapply(
    fit_memod, rep(3:5, 2), rep(c("m", "r"), each = 3)
    , MoreArgs = list(sbgrp_df = memod_fit_df, fixef_cols = c(ccd_keep_cols, "lang", "sg"),
                      year = 19, re_specs = "(1|tx_schoolid/sg)")
    , SIMPLIFY = FALSE
  )
  mods_idssubset_list <- mapply(
    fit_memod, rep(3:5, 2), rep(c("m", "r"), each = 3)
    , MoreArgs = list(sbgrp_df = memod_fit_df, fixef_cols = c(ccd_keep_cols, "lang", "sg"),
                      year = 19, ids_subset = unique(memod_fit_df$tx_schoolid[memod_fit_df$lang == "s" &
                                                                                !is.na(memod_fit_df$rs)]),
                      re_specs = "(1|tx_schoolid/sg)")
    , SIMPLIFY = FALSE
  )
  mods_langrfx_list <- mapply(
    fit_memod, rep(3:5, 2), rep(c("m", "r"), each = 3)
    , MoreArgs = list(sbgrp_df = memod_fit_df, fixef_cols = c(ccd_keep_cols, "lang", "sg"),
                      year = 19, re_specs = c("(1|tx_schoolid/sg)", "(1|tx_schoolid:lang)"))
    , SIMPLIFY = FALSE
  )
  names(mods_list) <- paste0("g", rep(3:5, 2), "_", rep(c("m", "r"), each = 3))
  # memods <- mapply(
  #   fit_memods,
  #   rep(c("m", "r"), 2),
  #   rep(c("e", "s"), each = 2),
  #   MoreArgs = list(sbgrp_df = memod_fit_df, W_col = "rs", grade_col = "grade",
  #                   Z_cols = ccd_keep_cols, id_col = "tx_schoolid", sbgrp_col = "sg",
  #                   year = 19),
  #   SIMPLIFY = FALSE
  # )
  # names(memods) <- paste(rep(c("m", "r"), 2), rep(c("e", "s"), each = 2), sep = "_")
  # memods <- setNames(
  #   Reduce(c, memods),
  #   paste(Reduce(c, mapply(names, memods, SIMPLIFY = FALSE)),
  #         mapply(function(ix) rep(names(memods)[ix], length(memods[ix])), seq_along(memods)),
  #         sep = "_")
  # )
  
  # make empirical bayes predictions for each grade/subject/language combo
  x_pred_df <- dplyr::distinct(
    memod_fit_df[memod_fit_df$join_yr == scrs_yr,
                 setdiff(colnames(memod_fit_df),
                         c("grade", "subj", "lang", "sg", "year.x", "year.y", "d", "rs"))]
  )
  x_pred_df <- dplyr::bind_cols(
    Reduce(dplyr::bind_rows, replicate(2 * length(sgs_keep), x_pred_df, simplify = FALSE)), # 2 for subj being either "m" or "r"
    lang = rep(rep(c("e", "s"), each = nrow(x_pred_df)), length(sgs_keep)),
    sg = rep(sgs_keep, each = nrow(x_pred_df) * 2)
  )
  xhats <- cbind(
    x_pred_df[, c("tx_schoolid", "sg", "lang")],
    mapply(
      function(mod, df) {
        fxfx <- fixef(mod)$cond
        randfx <- ranef(mod)$cond
        fxform <- delete.response(terms(formula(mod, fixed.only = TRUE, component = "cond")))
        # fixed effects
        drop(
          model.matrix(fxform
                       , data = model.frame(fxform, df, na.action = "na.pass"))[,names(fxfx)] %*%
            fxfx
        ) +
          # random effects
          rowSums(
            mapply(
              function(nm, rfx, df) {
                cols <- strsplit(nm, ":")[[1L]]
                grps <- do.call(paste, c(as.list(df[, cols, drop=FALSE]), list(sep = ":")))
                rfx_vec <- unname(setNames(rfx[[nm]][,1], row.names(rfx[[nm]]))[grps])
                replace(rfx_vec, is.na(rfx_vec), 0) # this mimics allow.new.levels = TRUE
              },
              names(mod$modelInfo$reTrms$cond$cnms),
              MoreArgs = list(rfx = randfx, df = df))
          )
      },
      mods_list,
      MoreArgs = list(df = x_pred_df),
      USE.NAMES = TRUE
    )
  )
  
  # get MLE pihats and match
  mles <- get_mle_logits(lr, xhats, scrs_yr, c("ethh", "ethb"), scr_cols[grepl("eth(h|b)", scr_cols)], K_mix_comps = 3)
  saveRDS(mles, file = "mle_ps_adsy.Rdata")
  saveRDS(lr_fit_df$Ti, file = "adsy_trt.Rdata")
  
  # get RC pihats and match
  rcs <- get_rc_logits(lr_form, lr_fit_df, xhats, scr_cols[grepl("eth(h|b)", scr_cols)], c("ethh", "ethb"), scrs_yr)
  saveRDS(rcs, file = "rc_ps_adsy.Rdata")

  # get naive pihats and match
  naive <- readRDS("naive_logits.Rdata")

  # combos: (0.7, 3), (0.7, 5), (0.7, 7), (1, 3), (1, 5), (1, 7)
  mle_match <- fullmatch_logits(lr_fit_df, "Ti", caliper_width = 0.5,
                                logits = setNames(mles, row.names(lr_fit_df)),
                                fullmatch_params = list(max.controls = 5),
                                grades_match, sch_type_match)
  rc_match <- fullmatch_logits(lr_fit_df, "Ti", caliper_width = 0.5,
                               logits = setNames(rcs, row.names(lr_fit_df)),
                               fullmatch_params = list(max.controls = 5),
                               grades_match, sch_type_match)
  naive_match <- fullmatch_logits(lr_fit_df, "Ti", caliper_width = 0.5,
                                  logits = setNames(naive, row.names(lr_fit_df)),
                                  fullmatch_params = list(max.controls = 5),
                                  grades_match, sch_type_match)
  
  sm <- Reduce(
    cbind,
    list(calc_balmetrics(mle_match, lr_fit_df, xhats, sgs, scrs_yr,
                         scr_cols[grepl("eth(h|b)", scr_cols)], "Ti",
                         scr_cols[grepl("eth(h|b)", scr_cols)]),
         calc_balmetrics(rc_match, lr_fit_df, xhats, sgs, scrs_yr,
                         scr_cols[grepl("eth(h|b)", scr_cols)],  "Ti",
                         scr_cols[grepl("eth(h|b)", scr_cols)]),
         calc_balmetrics(naive_match, lr_fit_df, xhats, sgs, scrs_yr,
                         scr_cols[grepl("eth(h|b)", scr_cols)], "Ti",
                         scr_cols[grepl("eth(h|b)", scr_cols)]),
         calc_balmetrics(rep("a", nrow(lr_fit_df)), lr_fit_df, xhats, sgs, scrs_yr,
                         scr_cols[grepl("eth(h|b)", scr_cols)], "Ti",
                         scr_cols[grepl("eth(h|b)", scr_cols)]))
  )
  colnames(sm) <- c("mle", "rc", "naive", "unmatched")
  n_matched_trts <- c(
    sum(!is.na(mle_match) & lr_fit_df[["Ti"]] == 1),
    sum(!is.na(rc_match) & lr_fit_df[["Ti"]] == 1),
    sum(!is.na(naive_match) & lr_fit_df[["Ti"]] == 1),
    sum(lr_fit_df[["Ti"]])
  )
  round(rbind(sm, n_matched_trts), 4)
  round(colMeans(sm[grepl("stdized", row.names(sm)),]), 4)

  # get matches and match statistics
  mle_match4 <- fullmatch_logits(lr_fit_df, "Ti", caliper_width = 1,
                                 logits = setNames(mles, row.names(lr_fit_df)),
                                 fullmatch_params = list(max.controls = 10),
                                 grades_match, sch_type_match)
  rc_match4 <- fullmatch_logits(lr_fit_df, "Ti", caliper_width = 1,
                                logits = setNames(rcs, row.names(lr_fit_df)),
                                fullmatch_params = list(max.controls = 10),
                                grades_match, sch_type_match)
  naive_match4 <- fullmatch_logits(lr_fit_df, "Ti", caliper_width = 1,
                                   logits = setNames(naive, row.names(lr_fit_df)),
                                   fullmatch_params = list(max.controls = 10),
                                   grades_match, sch_type_match)
  
  sm4 <- Reduce(
    cbind,
    list(calc_balmetrics(mle_match4, lr_fit_df, xhats, c("ethh", "ethb"), scrs_yr,
                         scr_cols[grepl("eth(h|b)", scr_cols)], "Ti",
                         scr_cols[grepl("eth(h|b)", scr_cols)]),
         calc_balmetrics(rc_match4, lr_fit_df, xhats, c("ethh", "ethb"), scrs_yr,
                         scr_cols[grepl("eth(h|b)", scr_cols)],  "Ti",
                         scr_cols[grepl("eth(h|b)", scr_cols)]),
         calc_balmetrics(naive_match4, lr_fit_df, xhats, c("ethh", "ethb"), scrs_yr,
                         scr_cols[grepl("eth(h|b)", scr_cols)], "Ti",
                         scr_cols[grepl("eth(h|b)", scr_cols)]),
         calc_balmetrics(rep("a", nrow(lr_fit_df)), lr_fit_df, xhats, c("ethh", "ethb"), scrs_yr,
                         scr_cols[grepl("eth(h|b)", scr_cols)], "Ti",
                         scr_cols[grepl("eth(h|b)", scr_cols)]))
  )
  colnames(sm4) <- c("mle", "rc", "naive", "unmatched")
  n_matched_trts <- c(
    sum(!is.na(mle_match4) & lr_fit_df[["Ti"]] == 1),
    sum(!is.na(rc_match4) & lr_fit_df[["Ti"]] == 1),
    sum(!is.na(naive_match4) & lr_fit_df[["Ti"]] == 1),
    sum(lr_fit_df[["Ti"]])
  )
  ess <- c(vapply(list(mle_match4, rc_match4, naive_match4)
                  , effectiveSampleSize
                  , numeric(1L)),
           NA_real_)
  round(rbind(sm4[grepl("stdized", row.names(sm4)),], n_matched_trts, ess), 4)
  round(colMeans(sm4[grepl("stdized", row.names(sm4)),]), 4)

  # estimate treatment effects
  mle_placebo_ests <- estimate_trt_effect(c("m", "r"), 3, 19, c("ethw", "etha")
                                          , get_covs(include_trt = TRUE)
                                          , schoolmatch = mle_match4
                                          , schls_to_estimate = treated_to_assess
                                          , memod_fit_df = memod_fit_df
                                          , lr_data = lr_fit_df)
  rc_placebo_ests <- estimate_trt_effect(c("m", "r"), 3, 19, c("ethw", "etha")
                                         , get_covs(include_trt = TRUE)
                                         , schoolmatch = rc_match4
                                         , schls_to_estimate = treated_to_assess
                                         , memod_fit_df = memod_fit_df
                                         , lr_data = lr_fit_df)
  naive_placebo_ests <- estimate_trt_effect(c("m", "r"), 3, 19, c("ethw", "etha")
                                            , get_covs(include_trt = TRUE)
                                            , schoolmatch = naive_match4
                                            , schls_to_estimate = treated_to_assess
                                            , memod_fit_df = memod_fit_df
                                            , lr_data = lr_fit_df)
  summary(mle_placebo_ests)
  summary(rc_placebo_ests)
  summary(naive_placebo_ests)
}