library("rstanmed")
library("batchtools")

dir <- paste0(getwd(), "/")
options(batchtools.progress = FALSE)


# Transportability violation and inflation factor -------------------------

reg_inf <- makeRegistry(paste0(dir, "inflation/"), make.default = TRUE)

nreps <- 200
batchMap(rstanmed:::run_bdf_replicate_tviol, 
         reg = reg_inf,
         args = data.table::CJ(seed = (1:nreps)*3, 
                               inflate_factor = c(1, 2.5, 5, 100),
                               n = 5000, u_ei = 1, am_intx = 1, 
                               yu_strength = 1.5, mu_strength = 1.5,
                               tviol = c(0, 2, 3),
                               n_ratio = 1,
                               result_type = "raw"))




# True NDERs --------------------------------------------------------------

loadResult(2, reg = reg_inf)$small_truth_nder
loadResult(3, reg = reg_inf)$small_truth_nder
loadResult(3, reg = reg_inf)$truth_nder


# Summary tables ----------------------------------------------------------

oc <- rstanmed:::combine_tviol_reg_results(registry = reg_inf)
oc$ci_cov_gc <- oc$ci_cov_gc * 100


tab_c <- aggregate(ci_cov_gc ~ inflate_factor + tviol, data = oc, FUN = mean)
tab_w <- aggregate(ci_width_gc ~ inflate_factor + tviol, data = oc, FUN = mean)


tab_c <- reshape2::dcast(data = tab_c, tviol ~ inflate_factor,
                         value.var = "ci_cov_gc")
tab_c$tviol <- factor(tab_c$tviol, levels = c(0, 2, 3), 
                      labels = c("Full transportability",
                                 "Potentially recoverable transportability violation",
                                 "Transportability violation in parameter not identified in main data"))

# Coverage table
gsub(pattern = "\\label{tab:}", fixed = TRUE, replacement = "\\label{tab:covif}", x = 
       gsub(pattern = "\\toprule", fixed = TRUE, x = 
              kableExtra::kable(tab_c, format = "latex",
                                col.names = c("Transportability scenario", colnames(tab_c)[-1]),
                                booktabs = TRUE, format.args = list(nsmall = 1),
                                caption = "Coverage percentages for 95\\% credible intervals for simulation-based (BDF-SIM) Bayesian data fusion estimators, calculated in 200 replicates with exposure-induced mediator-outcome confounding and $N_1 = N_2 = 5000$."),
            replacement = "\\toprule & \\multicolumn{4}{c}{Inflation factor $\\sigma$} \\\\ "))

# Width table
tab_w <- reshape2::dcast(data = tab_w, tviol ~ inflate_factor, 
                         value.var = "ci_width_gc")
tab_w$tviol <- factor(tab_w$tviol, levels = c(0, 2 ,3), 
                      labels = c("Full transportability",
                                 "Potentially recoverable transportability violation",
                                 "Transportability violation in parameter not identified in main data"))

gsub(pattern = "\\label{tab:}", fixed = TRUE, replacement = "\\label{tab:widif}", x = 
       gsub(pattern = "\\toprule", fixed = TRUE, x = 
              kableExtra::kable(tab_w, format = "latex",
                                col.names = c("Transportability scenario", colnames(tab_c)[-1]),
                                booktabs = TRUE, format.args = list(digits = 2),
                                caption = "Widths of 95\\% credible intervals for simulation-based (BDF-SIM) Bayesian data fusion estimators, calculated in 200 replicates with exposure-induced mediator-outcome confounding and $N_1 = N_2 = 5000$."),
            replacement = "\\toprule & \\multicolumn{4}{c}{Inflation factor $\\sigma$} \\\\ "))