#! /usr/bin/env Rscript

start_time <- Sys.time()

here::i_am("inst/sbc/make_reference_rankhist.R")
library(here)

## ensure that the current dev version of OncoBayes2 is build and
## loaded.  code to use if one wants the more conservative
## load_OB2_dev_install routine to be used
setwd(here())
if (system("make dev-install") != 0) {
        stop("make dev-install failed!")
}

setwd(here("inst", "sbc"))

# Always reload startup script when working interactively
options(sbc.job.startup.complete = FALSE)
source(here("inst", "sbc", "sbc_job_startup.R"))

library(clustermq)
library(knitr)
set.seed(453453)

## replications to use
S <- 1E4
##S <- 1E3
##S <- 2E3
##S <- 5E2
## S <- 10

examples_set <- names(example_models)
## temporarily reduced set of examples
##examples_set <- c("combo2_EX", "log2bayes_EXNEX")
##examples_set <- c("log2bayes_EXNEX")
##examples_set <- c("combo3_EXNEX", "log2bayes_EXNEX")

scenarios <- expand.grid(
        repl = 1:S,
        data_scenario = examples_set,
        stringsAsFactors = FALSE
)
scenarios <- cbind(job.id = 1:nrow(scenarios), scenarios)

num_simulations <- nrow(scenarios)

cat("Total number of jobs to dispatch:", num_simulations, "\n")

scheduler <- getOption("clustermq.scheduler")

if (is.null(scheduler)) {
        ## in this case we enable the multiprocess option to leverage local CPUs
        options(clustermq.scheduler = "multiprocess")
}

## used for debugging
## options(clustermq.scheduler="LOCAL")

scheduler <- getOption("clustermq.scheduler")

n_jobs <- 1
if (scheduler == "multiprocess") {
        ## on a local machine we only use as many CPUs as available
        n_jobs <- parallelly::availableCores(methods = "system")
}
if (scheduler %in% c("LSF", "SGE", "SLURM", "PBS", "Torque")) {
        ## on a queuing enabled backend, we use a lot more parallelism
        n_jobs <- 600
        ## n_jobs  <- 10
}

# Only use as many jobs as there are simulations at most
n_jobs <- min(n_jobs, num_simulations)

cat("Using clustermq backend", scheduler, "with", n_jobs, "concurrent jobs.\n")


RNGkind("L'Ecuyer-CMRG")
set.seed(56969839)
rng_seeds <- setup_lecuyer_seeds(.Random.seed, num_simulations)

scenarios$rng_seed <- rng_seeds

S_warmup <- max(floor(0.1 * S), 1)
S_final <- S - S_warmup

idx_warmup <- scenarios$repl <= S_warmup

scenarios_warmup <- scenarios[idx_warmup, ]
scenarios_main <- scenarios[!idx_warmup, ]

worker <- clustermq::workers(
        n_jobs,
        reuse = TRUE,
        template = list(
                walltime = 59,
                job_name = paste0("ob2_sbc[1-", n_jobs, "] #"),
                memory = 6 * 1024
        ),
        log_worker = FALSE
)

if (FALSE) {
        ## for debugging purposes
        res <- list()
        for (row in seq_len(nrow(scenarios_warmup))) {
                res[[row]] <- with(
                        scenarios_warmup,
                        run_sbc_case(
                                job.id[row],
                                repl[row],
                                data_scenario[row],
                                example_models,
                                rng_seeds
                        )
                )
        }
        res[[2]]
}

sim_result_warmup <- Q_rows(
        scenarios_warmup,
        run_sbc_case,
        const = list(example_models_cmq = example_models),
        n_jobs = n_jobs,
        workers = worker
)


## TODO: crew alternative? Should be implemented as we then get much better
## and easier debugging!!! Easiest debugging is via targets...
# if(FALSE) {
#
#     library(crew)
#
#     dir.create(here("log"), FALSE)
#     controller <- crew_controller_local(
#         name = "sbc",
#         workers = 10,
#         seconds_idle = 10,
#         local_log_directory=here("log")
#     )
#
#     sim_result_warmup <- controller$map(sbc_tools$run_sbc_case(job.id, repl, data_scenario, base_scenarios, seeds),
#                                         iterate=scenarios_warmup,
#                                         data=list(base_scenarios=sbc_tools$example_models, seeds=rng_seeds),
#                                         globals=c(as.list(sbc_tools), list(sbc_tools=sbc_tools)),
#                                         library=.libPaths(),
#                                         packages=pkg)
#
#     controller$error
#     controller$log
# }

assert_that(
        sum(idx_warmup) == length(sim_result_warmup),
        msg = "Check if all warmup simulations were processed."
)

#'
#' Collect warmup info
#'

warmup_info <- lapply(sim_result_warmup, function(run) {
        run[c("stepsize", "draws")]
})

scenarios_warmup <- transform(
        scenarios_warmup,
        warmup_job.id = 1:nrow(scenarios_warmup)
)

warmup_info_by_model <- lapply(
        split(scenarios_warmup$warmup_job.id, scenarios_warmup$data_scenario),
        function(mj) list(warmup_info = warmup_info[mj])
)

#'
#' Schedule remaining runs with available warmup info
#'

example_models_with_warmup <- modifyList(example_models, warmup_info_by_model)

sim_result_main <- Q_rows(
        scenarios_main,
        run_sbc_case,
        const = list(example_models_cmq = example_models_with_warmup),
        n_jobs = n_jobs,
        workers = worker
)

assert_that(
        sum(!idx_warmup) == length(sim_result_main),
        msg = "Check if all main simulations were processed."
)

## shut down cluster worker
## if proper cleanup is successful, cancel kill-on-exit
## see https://mschubert.github.io/clustermq/articles/technicaldocs.html

if (worker$cleanup()) {
        on.exit()
}

extract_results <- function(x) {
        c(
                rank = as.list(x$rank),
                list(
                        job.id = x$job.id,
                        time.running = x$time.running,
                        n_divergent = x$n_divergent,
                        min_Neff_bulk = x$min_Neff_bulk,
                        min_Neff_tail = x$min_Neff_tail,
                        max_Rhat = x$max_Rhat,
                        lp_ess_bulk = x$lp_ess_bulk,
                        lp_ess_tail = x$lp_ess_tail,
                        stepsize = x$stepsize,
                        accept_stat = x$accept_stat
                )
        )
}

calibration_data <- bind_rows(
        lapply(sim_result_warmup, extract_results),
        lapply(sim_result_main, extract_results)
) %>%
        arrange(job.id) %>%
        merge(scenarios, by = "job.id")

## check that indeed all jobs have finished
assert_that(nrow(calibration_data) == num_simulations)

## collect sampler diagnostics
sampler_diagnostics_summary <- calibration_data %>%
        group_by(data_scenario) %>%
        summarize(
                N = n(),
                total_divergent = sum(n_divergent),
                min_ess_bulk = min(min_Neff_bulk),
                min_ess_tail = min(min_Neff_tail),
                max_Rhat = max(max_Rhat),
                total_large_Rhat = sum(max_Rhat > 1.1),
                min_lp_ess_bulk = min(lp_ess_bulk),
                min_lp_ess_tail = min(lp_ess_tail)
        )


cat("\nSampler diagnostics:\n\n")
kable(sampler_diagnostics_summary, digits = 3)
cat("\n")

if (sum(sampler_diagnostics_summary$total_divergent) != 0) {
        warning("There were some divergent transitions!")
}
if (any(sampler_diagnostics_summary$max_Rhat > 1.1)) {
        warning("There were some parameters with large Rhat!")
}

#' Bin raw data as used in the analysis.
B <- 1024L / 2^5

rank_params <- names(calibration_data)[grepl(
        names(calibration_data),
        pattern = "rank"
)]

calibration_data_binned <- calibration_data %>%
        mutate(across(starts_with("rank"), function(x) {
                ceiling((x + 1) / (1024 / B) - 1)
        }))

## head(calibration_data_binned)

names(calibration_data_binned) <- gsub(
        names(calibration_data_binned),
        pattern = "rank[.]",
        replacement = ""
)

params <- gsub(rank_params, pattern = "rank[.]", replacement = "")

calibration_binned <- calibration_data_binned %>%
        dplyr::select(
                -job.id,
                -n_divergent,
                -min_Neff_bulk,
                -min_Neff_tail,
                -rng_seed
        ) %>%
        tidyr::gather(key = "param", value = "bin", -data_scenario) %>%
        group_by(data_scenario, param, bin) %>%
        tally() %>%
        right_join(
                expand.grid(
                        data_scenario = unique(
                                calibration_data_binned$data_scenario
                        ),
                        param = params,
                        bin = 0:(B - 1),
                        stringsAsFactors = FALSE
                ),
                c("data_scenario", "param", "bin")
        ) %>%
        replace_na(list(n = 0)) %>%
        arrange(data_scenario, param, bin) %>%
        spread(key = param, value = n)


#' Save as data.frame to avoid data.table dependency.
calibration_data <- as.data.frame(calibration_data)
calibration_binned <- as.data.frame(calibration_binned)

#' Further identification and verification data of run
git_hash <- system2("git", c("rev-parse", "HEAD"), stdout = TRUE)
created <- Sys.time()
created_str <- format(created, "%F %T %Z", tz = "UTC")

calibration <- list(
        ## raw = calibration_data, ## stop storing raw results, which are not needed for SBC reports
        data = calibration_binned,
        sampler_diagnostics = sampler_diagnostics_summary,
        S = S,
        B = B,
        git_hash = git_hash,
        created = created
)

saveRDS(calibration, file = here("inst", "sbc", "calibration.rds"))
## Store the large raw calibration data uncompressed: the corporate F5 WAF
## rejects the LFS upload PUT of the gzip-compressed variant (returning a
## block page as HTTP 200), which makes git-lfs report success while GitLab
## never receives the object and the pre-receive hook declines the push. The
## uncompressed serialization avoids the WAF signature and uploads cleanly.
saveRDS(calibration_data, file = here("inst", "sbc", "calibration_data.rds"), compress = FALSE)

library(tools)
md5 <- md5sum(here("inst", "sbc", "calibration.rds"))
cat(
        paste0(
                "Created:  ",
                created_str,
                "\ngit hash: ",
                git_hash,
                "\nMD5:      ",
                md5,
                "\n"
        ),
        file = here("inst", "sbc", "calibration.md5")
)


#'
#' Summarize execution time
#'

job_report <- calibration_data[c(
        "job.id",
        "repl",
        "data_scenario",
        "time.running"
)]
job_report$time.running <- job_report$time.running / 60
job_report$phase <- factor(idx_warmup, c(TRUE, FALSE), c("warmup", "main"))

runtime_by_problem_phase <- job_report %>%
        group_by(data_scenario, phase) %>%
        summarize(
                total = sum(time.running),
                mean = mean(time.running),
                max = max(time.running)
        ) %>%
        pivot_wider(
                id_cols = "data_scenario",
                names_from = "phase",
                values_from = c("total", "mean", "max")
        )

runtime_by_problem <- job_report %>%
        group_by(data_scenario) %>%
        summarize(
                total = sum(time.running),
                mean = mean(time.running),
                max = max(time.running)
        )

cat("Summary on job runtime on cluster:\n\n")

cat("\nRuntime by data scenario and phase:\n")
print(kable(runtime_by_problem_phase, digits = 2))

cat("\nRuntime by data scenario:\n")
print(kable(runtime_by_problem, digits = 2))

end_time <- Sys.time()

total_runtime <- difftime(end_time, start_time)
units(total_runtime) <- "mins"

cat("\n\nTotal runtime (min):", as.numeric(total_runtime), "\n\n\n")


#' Session info
print(sessionInfo())
