## -----------------------------------------------------------------------------
#| output: FALSE
library(whatifbandit)
library(dplyr)
library(tidyr)
library(ggplot2)
library(stringr)
library(forcats)
options(scipen = 9999)


## -----------------------------------------------------------------------------
data("tanf")
head(tanf)
glimpse(tanf)


## -----------------------------------------------------------------------------
set.seed(543645)
first_sim <- mab_from_rct(
  success ~ condition,
  data = tanf,
  algorithm = "ucb1",
  period_method = "batch",
  period_length = 1000,
  whole_experiment = TRUE
)


## -----------------------------------------------------------------------------
#| include: FALSE
load("tanf_sims.RData")


## -----------------------------------------------------------------------------
#| eval: FALSE
# future::plan("multisession", workers = 2)
# set.seed(53254)
# tanf_simulations <- mab_from_rct(
#   success ~ condition + block(service_center),
#   data = tanf,
#   algorithm = "thompson",
#   period_method = "date",
#   time_unit = "month",
#   period_length = 1,
#   delayed_feedback = TRUE,
#   assignment_date_col = letter_sent_date,
#   success_date_col = date_of_recert,
#   date_col = appt_date,
#   month_col = recert_month,
#   whole_experiment = FALSE,
#   random_assign_prop = 0.3,
#   r = 100,
#   keep_data = TRUE,
#   seed = TRUE
# )
# future::plan("sequential")


## -----------------------------------------------------------------------------
p <- matrix(
  c(0.20, 0.35, 0.5, 0.27, 0.33),
  ncol = 1,
  dimnames = list(
    c("control", "T1", "T2", "T3", "T4"),
    NULL
  )
)
print(p)

set.seed(123)
sim_from_scratch <- simulate_mab(
  n = 2000,
  t = 20,
  p = p,
  algorithm = "thompson",
  random_assign_prop = 0.3,
  contrasts = "all"
)


## -----------------------------------------------------------------------------
class(sim_from_scratch)
class(first_sim)
class(tanf_simulations)


## -----------------------------------------------------------------------------
head(first_sim$new_data)


## -----------------------------------------------------------------------------
head(sim_from_scratch$bandit$statistic)
head(sim_from_scratch$bandit$assignment_prob)
sim_from_scratch$bandit$assignment_quant

# Assignment probabilities over time
tanf_simulations$bandit$assignment_prob |>
  pivot_longer(
    cols = c(no_letter, open_appt, specific_appt),
    values_to = "p",
    names_to = "mab_condition"
  ) |>
  summarize(
    avg_p = mean(p),
    sd = sd(p),
    .by = c(mab_condition, period_number)
  ) |>
  mutate(
    lo = avg_p - 2 * sd,
    hi = avg_p + 2 * sd,
    mab_condition = str_to_title(str_replace(mab_condition, "_", " "))
  ) |>
  ggplot(aes(
    x = period_number,
    y = avg_p,
    color = mab_condition,
    fill = mab_condition
  )) +
  geom_ribbon(aes(ymax = hi, ymin = lo), alpha = 0.15, color = NA) +
  geom_line(linewidth = 0.8) +
  scale_y_continuous(labels = scales::percent) +
  labs(
    x = "Period",
    y = "Probability of Assignment",
    title = "Assignment Probability to Treatment
Arms Over Time",
    subtitle = "With 2 Standard Deviation Uncertainty Ribbon",
    color = "Treatment Arm",
    fill = "Treatment Arm"
  ) +
  theme_minimal() +
  theme(
    panel.grid.minor = element_blank(),
    panel.grid.major.y = element_blank(),
    plot.title = element_text(face = "bold"),
    plot.subtitle = element_text(face = "italic")
  )


## -----------------------------------------------------------------------------
head(sim_from_scratch$means)
# Plotting 95% CI using Normal distribution for AW-AIPW estimates
sim_from_scratch$means |>
  filter(estimator == "AW-AIPW") |>
  mutate(
    lo = mean + qnorm(0.025) * se,
    hi = mean + qnorm(0.975) * se,
    mab_condition = str_to_title(mab_condition),
    mab_condition = fct_reorder(mab_condition, mean)
  ) |>
  ggplot(aes(x = mean, y = mab_condition)) +
  geom_errorbar(
    aes(xmin = lo, xmax = hi),
    orientation = "y",
    width = 0.15,
    color = "grey40",
    linewidth = 0.6
  ) +
  geom_point(size = 2.5, color = "#2c3e50") +
  theme_minimal(base_size = 12) +
  labs(
    y = NULL,
    x = "Probability of Success",
    title = "Estimated Probability of Success By Arm",
    subtitle = "AW-AIPW Point Estimates with 95% Confidence Interval"
  ) +
  theme(
    panel.grid.major.y = element_blank(),
    plot.title = element_text(face = "bold"),
    plot.subtitle = element_text(face = "italic")
  )
# Simple test of no treatment effect between each pair
sim_from_scratch$contrasts |>
  filter(estimator == "AW-AIPW") |>
  mutate(
    z_stat = est / se,
    p_value = pnorm(abs(z_stat), lower.tail = FALSE) * 2
  ) |>
  select(arm1, arm2, est, se, z_stat, p_value)


## -----------------------------------------------------------------------------
# Selecting a random trial from the 100 resimulations
set.seed(12345)
i <- sample(100, 1)
bandit <- tanf_simulations$means |> filter(trial == i & estimator == "AW-AIPW")
rct <- estimatr::lm_robust(
  success ~ condition - 1,
  data = tanf,
  se_type = "HC2"
)

bandit
tibble(
  treatment = str_replace(names(rct$coefficients), "condition", ""),
  coefs = rct$coefficients,
  se = rct$std.error
)

bandit_assignments <- tanf_simulations$new_data |>
  unnest(data) |>
  filter(
    trial == i &
      mab_condition == "specific_appt"
  ) |>
  nrow()
regular_assignments <- tanf |> filter(condition == "specific_appt") |> nrow()

coef_diff <- abs(
  bandit$mean[bandit$mab_condition == "specific_appt"] -
    coefficients(rct)["conditionspecific_appt"]
)
se_diff <- abs(
  bandit$se[bandit$mab_condition == "specific_appt"] -
    rct$std.error["conditionspecific_appt"]
)


## -----------------------------------------------------------------------------
bandit |>
  bind_cols(
    bandit |>
      filter(mab_condition == "no_letter") |>
      select(control = mean, control_se = se)
  ) |>
  mutate(
    z_se = sqrt(se^2 + control_se^2),
    z_stat = (mean - control) / z_se,
    p_value = pnorm(abs(z_stat), lower.tail = FALSE) * 2
  ) |>
  filter(mab_condition != "no_letter") |>
  select(mab_condition, z = z_stat, se = z_se, p_value)


## -----------------------------------------------------------------------------
# joint_test parallelizes through the same options as the original sim
sim_from_scratch$config$parallel <- furrr::furrr_options(seed = TRUE)
set.seed(123)
boot <- joint_test(sim_from_scratch, method = "bootstrap", r = 100)
rand <- joint_test(sim_from_scratch, method = "randomization", r = 100)


## -----------------------------------------------------------------------------
boot$f_stat
boot$p_value
boot$method


## -----------------------------------------------------------------------------
#| include: FALSE
load("pure_sims.RData")


## -----------------------------------------------------------------------------
#| eval: FALSE
# set.seed(53245)
# pure_ucb_sims <- mab_from_rct(
#   success ~ condition,
#   data = tanf,
#   algorithm = "ucb1",
#   period_method = "individual",
#   keep_data = TRUE,
#   seed = TRUE,
#   r = 100,
#   whole_experiment = TRUE
# )
# pure_ts_sims <- update(pure_ucb_sims$config, algorithm = "thompson")


## -----------------------------------------------------------------------------
# RCT Mean Estimates
tanf_estimates <- coef(lm(success ~ condition - 1, tanf))
names(tanf_estimates) <- str_replace(names(tanf_estimates), "condition", "")
best <- names(tanf_estimates)[which.max(tanf_estimates)]

# Original RCT Data. Adding new columns to mimic output of MAB experimental data.
rct_regret <- tanf |>
  mutate(period_number = row_number(), mab_condition = condition)

# Selecting a random trial of the 100 previous simulations for regret analysis
set.seed(0934)
i <- sample.int(100, 1)

bandit_regrets <- lapply(list(pure_ucb_sims, pure_ts_sims), \(mab) {
  mab$new_data |> unnest(data) |> filter(trial == i)
})

lapply(
  list(
    "Static" = rct_regret,
    "UCB1" = bandit_regrets[[1]],
    "Thompson Sampling" = bandit_regrets[[2]]
  ),
  \(df) {
    df |>
      mutate(
        regret = tanf_estimates[[best]] - tanf_estimates[mab_condition]
      ) |>
      arrange(period_number) |>
      mutate(c_regret = cumsum(regret))
  }
) |>
  bind_rows(.id = "Algorithm") |>
  select(Algorithm, c_regret) |>
  mutate(x = row_number(), .by = Algorithm) |>
  ggplot(aes(x = x, y = c_regret, color = Algorithm)) +
  geom_line(linewidth = 0.8) +
  theme_minimal(base_size = 12) +
  labs(
    y = "Cumulative Regret",
    x = "Participants Assigned",
    title = "Cumulative Expected Regret By Adaptive Algorithm",
    subtitle = paste0(
      "Static Refers To Original RCT; Simulation ",
      i,
      " Selected"
    )
  ) +
  theme(
    panel.grid.major.y = element_blank(),
    plot.title = element_text(face = "bold"),
    plot.subtitle = element_text(face = "italic")
  )


## -----------------------------------------------------------------------------
tanf_benefits <- data.frame(
  hh_size = c(1, 2, 3, 4, 5, 6, 7, 8, 9, 10),
  benefit = c(504, 629, 803, 983, 1136, 1335, 1532, 1691, 1863, 2024)
)


## -----------------------------------------------------------------------------
true_success_n <- sum(tanf$success)
data_method <- lapply(
  list("Thompson Sampling" = pure_ts_sims, "UCB1" = pure_ucb_sims),
  \(mab) {
    mab$new_data |>
      unnest(data) |>
      summarize(
        add_success_n = sum(mab_success) - true_success_n,
        .by = trial
      ) |>
      summarize(mean_add_success_n = mean(add_success_n))
  }
) |>
  bind_rows(.id = "algorithm") |>
  pivot_wider(names_from = algorithm, values_from = mean_add_success_n) |>
  cbind(tanf_benefits) |>
  mutate(
    Welfare_TS = `Thompson Sampling` * benefit,
    Welfare_UCB1 = UCB1 * benefit
  ) |>
  select(hh_size, benefit, Welfare_TS, Welfare_UCB1)

data_method
lo <- min(data_method[1, c("Welfare_TS", "Welfare_UCB1")])
hi <- max(data_method[10, c("Welfare_TS", "Welfare_UCB1")])


## -----------------------------------------------------------------------------
expectation_method <- lapply(
  list("Thompson Sampling" = pure_ts_sims, "UCB1" = pure_ucb_sims),
  \(mab) {
    quants <- mab$bandit$assignment_quant[, -1] |> as.matrix()
    ests <- tanf_estimates[colnames(quants)]
    return(mean(quants %*% ests - true_success_n))
  }
) |>
  bind_rows(.id = "algorithm") |>
  cbind(tanf_benefits) |>
  mutate(
    Welfare_TS = `Thompson Sampling` * benefit,
    Welfare_UCB1 = UCB1 * benefit
  ) |>
  select(hh_size, benefit, Welfare_TS, Welfare_UCB1)
expectation_method