# =============================================================================
# Runner: I-SPY2 discrete simulation (population truth, recovery, top-k)
# =============================================================================
source("./Paper/Fun.R")
source("./Paper/Fun_Simulations.R")

ISPY2_PATH <- "./Paper/ISPY2-JCCELL-TableS2.xlsx"
OUT        <- "output"
INCLUDE_OR <- FALSE                # TRUE not supported under the current z_star 
                                   #fixed at population value 
N_GRID     <- c(8, 16, 32, 64, 128, 256)   # per-cell sizes for the top-k sweep
R_REPS     <- 1000
SEED_SIM       <- 1

dir.create(OUT, showWarnings = FALSE, recursive = TRUE)

# --- Module 1: population truth (flat Beta(1,1) posterior mean) ---------------
cells <- load_ispy2_cells(ISPY2_PATH, subtype = "Receptor Subtype")
truth <- build_truth_table(flat_truth_rates(cells), theta_or = 2)

# Population leapfrog table. With the flat-prior truth these population scores
# coincide with the real-data point estimates in Fun_Real.R, so this table also
# serves as a potential leapfrog table (not shown in the paper). 
lf_tab <- make_leapfrog_table(truth)
print(lf_tab)
export_leapfrog_latex(
  lf_tab,
  file    = file.path(OUT, "leapfrog_table.tex"),
  caption = "Population priority scores and induced global rankings for all candidate switches in the I-SPY2 calibration. Bold entries differ from the CATE ranking. Rows marked with $*$ are optimal and excluded from switching.",
  label   = "tab:leapfrog")

# --- Module 2: population scores and leapfrog counts --------------------------
pop_df <- get_population_scores(truth, include_or = INCLUDE_OR)
groups <- truth$groups

sw <- pop_df
sw$rank_cate_global <- rank(-sw$score_cate, ties.method = "first")
sw$rank_fr_global   <- rank(-sw$score_fr,   ties.method = "first")
sw$rank_ind_global  <- rank(-sw$score_ind,  ties.method = "first")

cat("Total candidate switches:", nrow(sw), "\n")
cat("Leapfrogs (Indep vs CATE):  ", sum(sw$rank_cate_global != sw$rank_ind_global),
    "/", nrow(sw), "\n")
cat("Leapfrogs (Frechet vs CATE):", sum(sw$rank_cate_global != sw$rank_fr_global),
    "/", nrow(sw), "\n")

# --- Module 3: finite-n engine (full grid; feeds top-k) ----------------------
sim_df <- run_sim_engine(groups, n_sizes = N_GRID, R = R_REPS,
                         theta_or = 2, include_or = INCLUDE_OR, seed = SEED_SIM)

# --- Module 4: rank-recovery raincloud ---------------------------------------
# Each estimated score vs its own population score; facets = common-n grid +
# Trial n (heterogeneous cell sizes). 
rec_df <- build_recovery_panels(groups, cells, pop_df,
                                n_grid = c(16, 64, 256), R = R_REPS,
                                include_or = INCLUDE_OR, theta_or = 2, seed = SEED_SIM)
plot_tau_recovery(rec_df, include_or = INCLUDE_OR,
                  output_dir = OUT, filename = "ispy2_tau_recovery.pdf")

# --- Module 5: score-vs-score heatmap ----------------------------------------
hm_df <- build_heatmap_panels(groups, cells, pop_df,
                              n_grid = c(16, 64, 256), R = R_REPS,
                              include_or = INCLUDE_OR, theta_or = 2, seed = SEED_SIM)
plot_tau_heatmap(hm_df, output_dir = OUT, filename = "ispy2_tau_heatmap.pdf")

# --- Module 6: top-k recovery (CATE, Indep, Frechet) -------------------------
set.seed(123)
topk_df <- compute_topk_recovery(sim_df, pop_df, ks = c(3, 6, 12),
                                 methods = c("cate", "ind", "fr"))
# Heterogeneous trial-size top-k (same statistic, actual I-SPY2 cell sizes)
sim_true  <- run_sim_engine_truen(groups, cells, R = R_REPS, theta_or = 2,
                                  include_or = INCLUDE_OR, seed = SEED_SIM+1)

set.seed(124)
topk_true <- compute_topk_recovery(sim_true, pop_df, ks = c(3, 6, 12),
                                   methods = c("cate", "ind", "fr"))
print(topk_true)

plot_topk_recovery(topk_df, topk_true = topk_true, realistic_n = range(cells$n_cell),
                   output_dir = OUT, filename = "ispy2_topk_recovery.pdf")
