source("./Paper/Fun.R")

# ============================================================================
# 1. DGP PARAMETERS
# Discrete Three-Group Data Generating Process — Multilevel Treatment
# Three arms (z = 1, 2, 3 in R indexing = arm 0, 1, 2 in paper)
# Baseline Policy: uniform across arms
# ============================================================================

groups_main <- list(
  G1 = list(gamma_vec = c(0.15, 0.65, 0.20), pol_vec = c(1/3, 1/3, 1/3), prob = 1/2),
  G2 = list(gamma_vec = c(0.55, 0.65, 0.95), pol_vec = c(1/3, 1/3, 1/3), prob = 1/2)
)

# ============================================================================
# 2. RANKING TABLE Introduction
# ============================================================================
# Generate table 
tab_main <- generate_leapfrog_table(groups_main, theta_or = c(5,7,10))
cat("=== Main table scores ===\n")
print(tab_main$scores[, c("group", "z_baseline", "z_star", "is_zstar",
                          "gamma_z", "cate",
                          "score_cate", "rank_cate",
                          "score_ind",  "rank_ind",
                          "score_fr",   "rank_fr")])
cat("\n=== Main table rankings ===\n")
print(tab_main$rankings)

scores   <- tab_main$scores
rankings <- tab_main$rankings

# ============================================================================
# 3. PARETO PLOT
# ============================================================================
theta_true_vals <- c(0, 1, 3)

pareto_main <- compute_and_plot_pareto_multilevel(
  tab             = tab_main$scores[!tab_main$scores$is_zstar, ],
  theta_true_vals = theta_true_vals,
  output_dir      = "./output",
  option          = "B"
)
