context('array')

Design <- createDesign(N = c(10, 20, 30))

Design5 <- expandDesign(Design, 5)

# iseed <- genSeeds()
iseed <- 554184288

Generate <- function(condition, fixed_objects) {
    dat <- with(condition, rnorm(N, 10, 5)) # distributed N(10, 5)
    dat
}

Analyse <- function(condition, dat, fixed_objects) {
    ret <- c(mean=mean(dat), median=median(dat)) # mean/median of sample data
    ret
}

Summarise <- function(condition, results, fixed_objects){
    colMeans(results)
}

test_that('array', {

    library(SimDesign)

    Design <- createDesign(N = c(10, 20, 30))

    Analyse.slow <- function(condition, dat, fixed_objects) {
        Sys.sleep(1)
        ret <- c(mean=mean(dat), median=median(dat)) # mean/median of sample data
        ret
    }

     # time test across conditions
    expect_error(suppressWarnings(runSimulation(design=Design, replications=4, generate=Generate,
                  analyse=Analyse.slow, summarise=Summarise,
                  control = list(max_time = "00:00:06", max_RAM = "4GB"),
                  verbose=FALSE)))
    files <- dir()
    file <- files[grepl('SIMDESIGN-TEMPFILE', files)]
    res <- SimRead(file)
    expect_true(is.null(res[[3]]))

    # resume from time crash
    res2 <- runSimulation(design=Design, replications=4, generate=Generate,
                          analyse=Analyse.slow, summarise=Summarise,
                          control = list(max_time = "00:00:20", max_RAM = "4GB"),
                          verbose=FALSE)
    expect_true(is(res2, 'SimDesign'))
    results <- SimResults(res2)
    # expect_true(nrow(results) == 4*nrow(Design))

    # define initial seed (do this only once to keep it constant!)
    # iseed <- genSeeds()
    iseed <- 554184288

    ### On cluster submission, the active array ID is obtained via getArrayID(),
    ###   and therefore should be used in real SLURM submissions
    expect_warning(arrayID <- getArrayID(type = 'slurm'))

    # However, for the following example array ID is set to first row only
    arrayID <- 1L

    # run the simulation (results not caught on job submission, only files saved)
    res <- runArraySimulation(design=Design, replications=50,
                              generate=Generate, analyse=Analyse,
                              summarise=Summarise, arrayID=arrayID,
                              iseed=iseed, filename='mysim',
                              control=list(store_Random.seeds=TRUE), verbose=FALSE)
    results <- SimExtract(res, what='results') # condition and replication count stored
    expect_equal(results$condition[1], 1)
    expect_equal(results$arrayID[1], 1)
    randseeds <- SimExtract(res, what='Random.seeds')
    expect_true(all(dim(randseeds[[1L]]) == c(50,7)))

    SimClean('mysim-1.rds')

    # serial iseed must be byte-for-byte reproducible (L'Ecuyer-CMRG seed
    # applied via set_seed(); see GitHub bug re: .GlobalEnv assignment)
    res_s1 <- runArraySimulation(design=Design, replications=20,
                                 generate=Generate, analyse=Analyse,
                                 summarise=Summarise, arrayID=1L,
                                 iseed=iseed, filename='serialseed',
                                 parallel=FALSE, verbose=FALSE)
    SimClean('serialseed-1.rds')
    res_s2 <- runArraySimulation(design=Design, replications=20,
                                 generate=Generate, analyse=Analyse,
                                 summarise=Summarise, arrayID=1L,
                                 iseed=iseed, filename='serialseed',
                                 parallel=FALSE, verbose=FALSE)
    expect_identical(SimExtract(res_s1, what='results'),
                     SimExtract(res_s2, what='results'))
    SimClean('serialseed-1.rds')

})

test_that('split array', {



    ########################
    # Same submission job as above, however split the replications over multiple
    # evaluations and combine when complete

    # arrayID <- getArrayID(type = 'slurm')
    arrayID <- 14L

    # run the simulation (replications reduced per row, but same in total)
    runArraySimulation(design=Design5, replications=10,
                       generate=Generate, analyse=Analyse,
                       summarise=Summarise, iseed=iseed,
                       filename='mylongsim', arrayID=arrayID, verbose=FALSE)

    res <- readRDS('mylongsim-14.rds')
    results <- SimResults(res) # condition and replication count stored
    expect_equal(results$condition[1], 3)
    expect_equal(results$arrayID[1], 14)
    randseeds <- SimExtract(res, what='Random.seeds')
    expect_true(is.null(randseeds))

    SimClean('mylongsim-14.rds')

    Analyse_big <- function(condition, dat, fixed_objects) {
        ret <- runif(1e4)
        names(ret) <- paste0('x', 1:length(ret))
        # object.size(ret) |> format('MB')
        ret
    }

    tmp <- gc()
    runArraySimulation(design=Design5, replications=1000,
                       generate=Generate, analyse=Analyse_big,
                       summarise=Summarise, iseed=iseed,
                       filename='mylongsim', arrayID=arrayID,
                       control = list(max_RAM="500MB"), verbose=FALSE)
    res <- readRDS("mylongsim-14.rds")
    expect_true(res$REPLICATIONS == 1000L)
    SimClean('mylongsim-14.rds')

    ###
    # Emulate distribution to nrow(Design5) = 15 independent job arrays
    sapply(1:nrow(Design5), \(arrayID)
           runArraySimulation(design=Design5, replications=10,
                              generate=Generate, analyse=Analyse,
                              summarise=Summarise, iseed=iseed, arrayID=arrayID,
                              dirname='sim', filename='condition',   # saved to 'sim/condition-#.rds'
                              control = list(max_time = "04:00:00", max_RAM = "4GB"),
                              verbose=FALSE)) |> invisible()

    #  If necessary, conditions above will manually terminate before
    #  4 hours, returning any successfully completed results before the HPC
    #  session times out (provided shell specified more than 4 hours)

    # list saved files
    files <- dir('sim/')
    expect_true(all(files %in% paste0('condition-', 1:nrow(Design5), '.rds')))

    setwd('sim')
    condition14 <- readRDS('condition-14.rds')
    results <- SimExtract(condition14, 'results')
    expect_equal(results$condition[1], 3)
    expect_equal(results$arrayID[1], 14)

    # aggregate simulation results into single file
    final <- SimCollect(files=dir())
    so <- summary(final)
    expect_equal(so$ncores, 15L)
    results <- SimResults(final)

    expect_equal(final$REPLICATIONS, c(50, 50, 50))
    expect_equal(dim(results), c(150, 5))

    setwd('..')
    SimClean(dirs='sim/')

})

test_that('lists', {

    # list return
    Analyse_list <- function(condition, dat, fixed_objects) {
        list(a=1:2, b=3:4)
    }

    Summarise_list <- function(condition, results, fixed_objects) {
        42
    }

    arrayID <- 1

    runArraySimulation(design=Design5, replications=10,
                       generate=Generate, analyse=Analyse_list,
                       summarise=Summarise_list, iseed=iseed, arrayID=arrayID,
                       dirname='sim', filename='condition', verbose=FALSE) |> invisible()
    res <- readRDS("sim/condition-1.rds")
    results <- SimExtract(res, 'results')
    expect_true(is.list(results))
    SimClean(dirs="sim/")

    # emulate the arrayID distribution, storing all results in a 'sim/' folder
    dir.create('sim/')

    # Emulate distribution to nrow(Design5) = 15 independent job arrays
    sapply(1:nrow(Design5), \(arrayID)
           runArraySimulation(design=Design5, replications=10,
                              generate=Generate, analyse=Analyse_list,
                              summarise=Summarise_list, iseed=iseed, arrayID=arrayID,
                              dirname='sim', filename='condition', verbose=FALSE)) |> invisible()

    files <- dir('sim/')
    expect_true(all(files %in% paste0('condition-', 1:nrow(Design5), '.rds')))

    setwd('sim')
    condition14 <- readRDS('condition-14.rds')
    results <- SimExtract(condition14, 'results')
    expect_equal(results[[1]]$condition, 3)
    expect_equal(results[[1]]$arrayID, 14)

    # aggregate simulation results into single file
    final <- SimCollect(files=dir())
    so <- summary(final)
    expect_equal(so$ncores, 15L)
    results <- SimResults(final)

    expect_equal(final$REPLICATIONS, c(50, 50, 50))
    expect_equal(length(results), 150)

    setwd('..')
    SimClean(dirs='sim/')

    #####################

    # summary list

    Summarise_list <- function(condition, results, fixed_objects) {
        list(a=42, b=rnorm(5))
    }

    # Emulate distribution to nrow(Design5) = 15 independent job arrays
    sapply(1:nrow(Design5), \(arrayID)
           runArraySimulation(design=Design5, replications=10,
                              generate=Generate, analyse=Analyse,
                              summarise=Summarise_list, iseed=iseed, arrayID=arrayID,
                              dirname='sim', filename='condition', verbose=FALSE)) |> invisible()

    files <- dir('sim/')
    expect_true(all(files %in% paste0('condition-', 1:nrow(Design5), '.rds')))

    setwd('sim')
    final <- SimCollect(files=dir())

    so <- summary(final)
    expect_equal(so$ncores, 15L)
    results <- SimResults(final)

    expect_equal(final$REPLICATIONS, c(50, 50, 50))
    expect_equal(nrow(results), 150)

    setwd('..')
    SimClean(dirs='sim/')



    #############
    # analyse returns a list
    Analyse <- function(condition, dat, fixed_objects) {
        ret <- list(mean=mean(dat), median=median(dat), stuff=rnorm(5))
        ret
    }

    Summarise <- function(condition, results, fixed_objects){
        means <- do.call(c, sapply(results, \(x) x['mean']))
        medians <- do.call(c, sapply(results, \(x) x['median']))
        c(means=mean(means), medians=mean(medians))
    }

    sapply(1:nrow(Design5), \(arrayID)
           runArraySimulation(design=Design5, replications=10,
                              generate=Generate, analyse=Analyse,
                              summarise=Summarise, iseed=iseed, arrayID=arrayID,
                              dirname='sim', filename='condition', verbose=FALSE)) |> invisible()

    files <- dir('sim/')
    expect_true(all(files %in% paste0('condition-', 1:nrow(Design5), '.rds')))

    setwd('sim')
    final <- SimCollect(files=dir())

    so <- summary(final)
    expect_equal(so$ncores, 15L)
    results <- SimResults(final)

    expect_equal(final$REPLICATIONS, c(50, 50, 50))
    expect_equal(length(results), 150)

    setwd('..')
    SimClean(dirs='sim/')

})

