# From Readme.md: https://earth.bsc.es/gitlab/es/startR

# ------------------------------------
# 1. Declaration of data sources
# ------------------------------------
library(startR)
rm(list = ls())
repos <- '/esarchive/exp/ecmwf/system5_m1/6hourly/$var$/$var$_$sdate$.nc'
data <- Start(dat = repos,
              var = 'tas',
              sdate = '20180101',
              ensemble = 'all',
              time = indices(1:2),
              latitude = 'all', #indices(1:40),
              longitude = indices(1:40),
              return_vars = list(latitude = 'dat', longitude = 'dat', time = 'sdate'),
              # retrieve = TRUE)
              retrieve = FALSE)
#-------------------------------------------------------------------

# ------------------------------------
# 2. Declaration of the workflow
# ------------------------------------

fun <- function(x) {
  # Expected inputs:
  #   x: array with dimensions ('ensemble', 'time')
  apply(x + 1, 2, mean)
}

step <- Step(fun,
             target_dims = c('ensemble', 'time'),
             output_dims = c('time'))

wf <- AddStep(data, step)

# ------------------------------------------------
# 2. Declaration of the HPC platform and execution
# ------------------------------------------------
res <- Compute(wf,
            chunks = list(latitude = 2,
                          longitude = 2))

res9 <- Compute(wf,
               chunks = list(latitude = 2,
                             longitude = 2),
               threads_load = 1,
               threads_compute = 2,
               cluster = list(queue_host = 'cte-power',
                              queue_type = 'slurm',
                              # temp_dir = '/gpfs/scratch/bsc32/bsc32339/test-startR/',
                              lib_dir = '/gpfs/projects/bsc32/share/R_libs/3.5/',
                              r_module = 'R/3.5.0-foss-2018b',
                              cores_per_job = 2,
                              max_jobs = 4,
                              job_wallclock = '00:10:00',
                              # extra_queue_params = list('#SBATCH --mem-per-cpu=3000'),
                              extra_queue_params = list('#SBATCH --qos=bsc_es'),
                              bidirectional = FALSE,
                              polling_period = 100),
               ecflow_suite_dir = '/esarchive/scratch/nperez/ecflow/')


resFat <- Compute(wf,
               chunks = list(latitude = 2, longitude = 2),
               threads_load = 1,
               threads_compute = 4,
               cluster = list(queue_host = 'bsceslogin01.bsc.es',
                              queue_type = 'slurm',
                              # temp_dir = '/home/Earth/nperez/startR_hpc/',
                              cores_per_job = 2,
                              job_wallclock = '00:10:00',
                              max_jobs = 4,
                              bidirectional = TRUE
                             ),
               ecflow_suite_dir = '/home/Earth/nperez/')
resmn <- Compute(wf,
               chunks = list(latitude = 2,
                             longitude = 2),
               threads_load = 1,
               threads_compute = 2,
               cluster = list(queue_host = 'mn1.bsc.es',
                              queue_type = 'slurm',
                              data_dir = '/gpfs/projects/bsc32/share/startR_data_repos/',
                              # temp_dir = '/gpfs/scratch/pr1efe00/pr1efe03/startR_tests/',
                              lib_dir = '/gpfs/projects/bsc32/share/R_libs/3.4/',
                              r_module = 'R/3.4.0',
                              cores_per_job = 2,
                              job_wallclock = '00:10:00',
                              max_jobs = 4,
                              extra_queue_params = list('#SBATCH --qos=prace'),
                              bidirectional = FALSE,
                              polling_period = 10,
                              special_setup = 'marenostrum4'
                             ),
               ecflow_suite_dir = '/home/Earth/nperez/',
               ecflow_server = NULL,
               silent = FALSE,
               debug = FALSE,
               wait = TRUE)


