### Copyright (C) 2018-2023 Marco De Lucia, Max Luebke (GFZ Potsdam)
###
### POET is free software; you can redistribute it and/or modify it under the
### terms of the GNU General Public License as published by the Free Software
### Foundation; either version 2 of the License, or (at your option) any later
### version.
###
### POET is distributed in the hope that it will be useful, but WITHOUT ANY
### WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
### A PARTICULAR PURPOSE. See the GNU General Public License for more details.
###
### You should have received a copy of the GNU General Public License along with
### this program; if not, write to the Free Software Foundation, Inc., 51
### Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.

master_init <- function(setup, out_dir, init_field) {
    ## Setup the directory where we will store the results
    if (!dir.exists(out_dir)) {
        dir.create(out_dir)
        msgm("created directory ", out_dir)
    } else {
        msgm("dir ", out_dir, " already exists, I will overwrite!")
    }
    if (is.null(setup$store_result)) {
        msgm("store_result doesn't exist!")
    } else {
        msgm("store_result is ", setup$store_result)
    }

    setup$iter <- 1
    setup$timesteps <- setup$timesteps
    setup$maxiter <- length(setup$timesteps)
    setup$iterations <- setup$maxiter
    setup$simulation_time <- 0

    dgts <- as.integer(ceiling(log10(setup$maxiter)))
    ## string format to use in sprintf
    fmt <- paste0("%0", dgts, "d")

    if (is.null(setup[["store_result"]])) {
        setup$store_result <- TRUE
    }

    if (setup$store_result) {
        init_field_out <- paste0(out_dir, "/iter_", sprintf(fmt = fmt, 0), ".", setup$out_ext)
        init_field <- data.frame(init_field, check.names = FALSE)
        SaveRObj(x = init_field, path = init_field_out)
        msgm("Stored initial field in ", init_field_out)
        if (is.null(setup[["out_save"]])) {
            setup$out_save <- seq(1, setup$iterations)
        }
    }

    setup$out_dir <- out_dir

    return(setup)
}

## This function, called only by master, stores on disk the last
## calculated time step if store_result is TRUE and increments the
## iteration counter
master_iteration_end <- function(setup, state_T, state_C) {
    iter <- setup$iter
    # print(iter)
    ## max digits for iterations
    dgts <- as.integer(ceiling(log10(setup$maxiter + 1)))
    ## string format to use in sprintf
    fmt <- paste0("%0", dgts, "d")

    ## Write on disk state_T and state_C after every iteration
    ## comprised in setup$out_save
    if (setup$store_result) {
        if (iter %in% setup$out_save) {
            nameout <- paste0(setup$out_dir, "/iter_", sprintf(fmt = fmt, iter), ".", setup$out_ext)
            state_T <- data.frame(state_T, check.names = FALSE)
            state_C <- data.frame(state_C, check.names = FALSE)

            ai_surrogate_info <- list(
                prediction_time = if (exists("ai_prediction_time")) as.integer(ai_prediction_time) else NULL,
                training_time = if (exists("ai_training_time")) as.integer(ai_training_time) else NULL,
                valid_predictions = if (exists("validity_vector")) validity_vector else NULL
            )

            SaveRObj(x = list(
                T = state_T,
                C = state_C,
                simtime = as.integer(setup$simulation_time),
                totaltime = as.integer(totaltime),
                ai_surrogate_info = ai_surrogate_info
            ), path = nameout)
            msgm("results stored in <", nameout, ">")
        }
    }
    ## Add last time step to simulation time
    setup$simulation_time <- setup$simulation_time + setup$timesteps[iter]

    ## msgm("done iteration", iter, "/", length(setup$timesteps))
    setup$iter <- setup$iter + 1
    return(setup)
}


## Attach the name of the calling function to the message displayed on
## R's stdout
msgm <- function(...) {
    prefix <- paste0("R: ")
    cat(paste(prefix, ..., "\n"))
    invisible()
}


GetWorkPackageSizesVector <- function(n_packages, package_size, len) {
    ids <- rep(1:n_packages, times = package_size, each = 1)[1:len]
    return(as.integer(table(ids)))
}


## Handler to read R objs from binary files using either builtin
## readRDS(), qs::qread() or qs2::qs_read() based on file extension
ReadRObj <- function(path) {
    ## code borrowed from tools::file_ext()
    pos <- regexpr("\\.([[:alnum:]]+)$", path)
    extension <- ifelse(pos > -1L, substring(path, pos + 1L), "")

    switch(extension,
        rds = readRDS(path),
        qs  = qs::qread(path),
        qs2 = qs2::qs_read(path)
    )
}

## Handler to store R objs to binary files using either builtin
## saveRDS() or qs::qsave() based on file extension
SaveRObj <- function(x, path) {
    ## msgm("Storing to", path)
    ## code borrowed from tools::file_ext()
    pos <- regexpr("\\.([[:alnum:]]+)$", path)
    extension <- ifelse(pos > -1L, substring(path, pos + 1L), "")

    switch(extension,
        rds = saveRDS(object = x, file = path),
        qs  = qs::qsave(x = x, file = path),
        qs2 = qs2::qs_save(object = x, file = path)
    )
}


######## Old relic code

## ## Function called by master R process to store on disk all relevant
## ## parameters for the simulation
## StoreSetup <- function(setup, filesim, out_dir) {
##     to_store <- vector(mode = "list", length = 4)
##     ## names(to_store) <- c("Sim", "Flow", "Transport", "Chemistry", "DHT")
##     names(to_store) <- c("Sim", "Transport", "DHT", "Cmdline")

##     ## read the setup R file, which is sourced in kin.cpp
##     tmpbuff <- file(filesim, "r")
##     setupfile <- readLines(tmpbuff)
##     close.connection(tmpbuff)

##     to_store$Sim <- setupfile

##     ## to_store$Flow <- list(
##     ##     snapshots  = setup$snapshots,
##     ##     gridfile   = setup$gridfile,
##     ##     phase      = setup$phase,
##     ##     density    = setup$density,
##     ##     dt_differ  = setup$dt_differ,
##     ##     prolong    = setup$prolong,
##     ##     maxiter    = setup$maxiter,
##     ##     saved_iter = setup$iter_output,
##     ##     out_save   = setup$out_save )

##     to_store$Transport <- setup$diffusion

##     ## to_store$Chemistry <- list(
##     ##    nprocs   = n_procs,
##     ##    wp_size  = work_package_size,
##     ##    base     = setup$base,
##     ##    first    = setup$first,
##     ##    init     = setup$initsim,
##     ##    db       = db,
##     ##    kin      = setup$kin,
##     ##    ann      = setup$ann)

##     if (dht_enabled) {
##         to_store$DHT <- list(
##             enabled   = dht_enabled,
##             log       = dht_log
##             ## signif    = dht_final_signif,
##             ## proptype  = dht_final_proptype
##         )
##     } else {
##         to_store$DHT <- FALSE
##     }

##     if (dht_enabled) {
##         to_store$DHT <- list(
##             enabled   = dht_enabled,
##             log       = dht_log
##             # signif    = dht_final_signif,
##             # proptype  = dht_final_proptype
##         )
##     } else {
##         to_store$DHT <- FALSE
##     }

##     saveRDS(to_store, file = paste0(fileout, "/setup.rds"))
##     msgm("initialization stored in ", paste0(fileout, "/setup.rds"))
## }