# single model:
          fit(workflow, data)

# multiple models:
fit_resamples(workflow, resamples(data))

library(tidymodels)
library(causalworkshop)
library(propensity)

net_data <- net_data %>% mutate(net = factor(net, levels = c("TRUE", "FALSE")))

outcome_wf <- 
  workflow(
    malaria_risk ~ net,
    linear_reg()
  ) %>%
  add_case_weights(wts)

propensity_wf <- 
  workflow(
    net ~ income + health + temperature,
    logistic_reg()
  )

net_data_wts <-
  fit(propensity_wf, net_data) %>%
  augment(net_data) %>%
  mutate(
    wts = wt_ate(.pred_TRUE, net, .treated = "TRUE")
  )

results <- 
  outcome_wf %>%
  fit(
    data = net_data_wts %>% mutate(wts = importance_weights(wts))
  ) %>%
  tidy()

results <-
  fit_resamples(
    propensity_wf,
    resamples = bootstraps(net_data),
    control = control_resamples(extract = identity)
  ) %>%
  weight_propensity(wt_ate) %>%
  fit_resamples(
    outcome_wf,
    resamples = .,
    control = control_resamples(extract = tidy)
  )

# a function that takes in a resample fit object and outputs a modified
# version of that object where the training data underlying each rsplit
# is augmented with propensity weights for each element of the analysis set. 
# this serves as the "bridge" between two calls to
# `fit_resamples()` (or `tune_*()`) in a causal workflow, the
# first being for the propensity model and the second for the outcome model.
# `tune_results` must have been executed with option `extract = identity`.
weight_propensity <- function(tune_results, wt_fn) {
  for (resample in seq_along(tune_results$splits)) {
    tune_results$splits[[resample]] <- 
      augment_split(
        tune_results$splits[[resample]],
        tune_results$.extracts[[resample]]$.extracts[[1]], 
        wt_fn = wt_fn,
        outcome_name = outcome_names(tune_results)
      )
  }
  
  tibble::new_tibble(
    tune_results[, c("splits", "id")], 
    !!!attr(tune_results, "rset_info")$att,
    class = c(attr(tune_results, "rset_info")$att$class, "rset")
  )
}

augment_split <- function(split, workflow, wt_fn, outcome_name) {
  d <- analysis(split)
  d <- vctrs::vec_cbind(d, predict(workflow, d, type = "prob"))
  d <- vctrs::vec_slice(d, !duplicated(d$id))

  model_fit <- extract_fit_parsnip(workflow)
  lvls <- model_fit$lvl
  event_lvl <- lvls[1]
  preds <- d[[paste0(".pred_", event_lvl)]]

  split[["data"]][d$id, "wts"] <- 
    importance_weights(wt_fn(preds, d[[outcome_name]], .treated = event_lvl))
  
  split
}

result <- 
  fit(
    propensity_wf, 
    net_data
  ) %>%
  weight_propensity(net_data, wt_ate, .) %>%
  fit(
    outcome_wf %>% add_case_weights(wts),
    data = .
  ) %>%
  tidy()