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

# ------------------------------------
# 1. Declaration of data sources
# ------------------------------------
library(startR)

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)
#-------------------------------------------------------------------
# checking output:
class(data) # retrieve TRUE: startR_array if retrieve FALSE: startR_cube
str(data)  #retrieve TRUE/FALSE: similar a big list of lists
names(data) # retrieve TRUE: NULL / retrieve FALSE: list of names with the firs one empty
data[[1]] # retrieve TRUE: 251.8854 / retrieve FALSE: Start
data$latitude # retrieve TRUE: it returns correctly an error / retrieve FALSE: all
attributes(data)$Variables$common$tas$dim[[1]]$name # retrieve TRUE longitude / FALSE NULL
attributes(data)$Variables$common$tas$dim[[1]]$vals
attributes(data)$Variables$common$tas$dim[[2]]$name # latitude /NULL
attributes(data)$Variables$common$tas$dim[[2]]$vals


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

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

#-------------------------------------------------------------------
mean(data[1,1,1, ,1,1,1] + 1) # result 252.8138
apply(data[1, 1, 1, , , 1, 1] + 1, 2, mean) # result 252.8138 252.4913
#-------------------------------------------------------------------
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(workflow = wf, 
            chunks = list(latitude = 2,
                          longitude = 2))

#-------------------------------------------------------------------
dim(res$output) # result time 2  dat  1 var 1  sdate 1  latitude 181 longitude 40
res$output[1 : 2, 1, 1, 1, 1, 1] # result  252.8138 252.4913
#-------------------------------------------------------------------

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/')


# ------------------------------------------------------
res9$output1[, 1, 1, 1, 1, 1] # result 252.8138 252.4913
# ------------------------------------------------------- 

# --------------------------------------------------------------------------------------
# Deborah code: with different data 
# -------------------------------------------------------------------------------------
rm(list = ls())
repos <- '/esarchive/exp/ecmwf/system5_m1/monthly_mean/$var$_f6h/$var$_$sdate$.nc'
data <- Start(dat = repos,
              var = 'tas',
              sdate = c('20170101', '20180101'),
              ensemble = indices(1 : 20),
              time = 'all',
              latitude = 'all', #indices(1:40),
              longitude = indices(1 : 40),
              return_vars = list(latitude = 'dat', longitude = 'dat', time = 'sdate'),
              #retrieve = TRUE)
              retrieve = FALSE)
source("./git/s2dverification/R/Season.R")
library(multiApply)
fun_spring <- function(x) {
    y <- Season(x,  monini = 1, moninf = 3, monsup = 5)
    return(y)
}
step <- Step(fun = fun_spring,
             target_dims = c('time'),
             output_dims = c('time'))#,

wf <- AddStep(data, step)

# Locally
library(tictoc)
tic()
res <- Compute(wf,
            chunks = list(ensemble = 2,
                          sdate = 2))
toc()
# To be checked:
res9 <- Compute(wf,
               chunks = list(ensemble = 2, sdate=2),
               threads_load = 1,
               threads_compute = 4,
               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 = 'startR/0.1.2-foss-2018b-R-3.5.0',
                              cores_per_job = 2,
                              job_wallclock = '00:30:00',
                              max_jobs = 4,
                              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/') #,
               #wait = TRUE)

resFat <- Compute(wf,
               chunks = list(ensemble = 2, sdate=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/')


res <- 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)


# ---------------------------------------
# Testing variable environment on Compute:
# ---------------------------------------
funp <- function(x) {
          
          return(c(ls(), class(x), mode(x), names(attributes(x)), dim(x)))
        }
step <- Step(fun = funp,
             target_dims = 'time',
             output_dims = c('class'),
             use_attributes = list(a = "B"))#, 'mode', 'attr'))#,
wf <- AddStep(data, step)
tic()
res <- Compute(wf,
            chunks = list(ensemble = 2,
                          sdate = 2))
toc()   
res$output1[,1,1,1,1,1,1] # result [1] "1"       "array"   "numeric" "dim"     "7"      


# To be checked:
res9 <- Compute(wf,
               chunks = list(ensemble = 2, sdate = 2),
               threads_load = 1,
               threads_compute = 4,
               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 = 'startR/0.1.2-foss-2018b-R-3.5.0',
                              cores_per_job = 2,
                              job_wallclock = '00:30:00',
                              max_jobs = 4,
                              #extra_queue_params = list('#SBATCH --mem-per-cpu=3000'),
                              extra_queue_params = list('#SBATCH --qos=bsc_es'),
                              bidirectional = FALSE,
                              # polling_period = 10),
               ecflow_suite_dir = '/esarchive/scratch/nperez/ecFlow/') #,
               #wait = TRUE)

# -------------------------
# Regrid sript and fun
# ------------------------
fun_two <- function(x) {
    y <- Season(x,  monini = 1, moninf = 3, monsup = 5)
    SelBox(y, lon = 
    return(y)
}


# ------------------------ Extra checks ----------------

library(multiApply)
res <- Apply(list(a), fun = fun_spring, target_dims = c('time'), output_dims = 'time')

# Check Apply:
a <- 1 : (2 * 59 * 10 * 12)
dim(a) <- c(dat = 2, time = 59, lat = 10, lon = 12)
res <- Apply(list(a), fun = fun_spring, target_dims = c('time'), output_dims = 'time')
a <- 1 : (2 * 7 * 10 * 12)
dim(a) <- c(dat = 2, time = 7, lat = 10, lon = 12)
res <- Apply(list(a), fun = fun_spring, target_dims = c('time'), output_dims = 'time')
res <- Apply(list(a), fun = fun_spring, target_dims = c('time'))

Season(1 : 59, posdim = 1, monini = 1, moninf = 3, monsup = 5)# result 4 16 28 40 52

library(multiApply)
fun <- function(x) {
    pos <- cbind(seq(3, length(x), <2), seq(4, length(x), 12), seq(5, length(x), 12))
    apply(pos, 1, function(y) {mean(x[y])})
}
fun(x = 1 : 59)
res_time <- Apply(list(a), fun = fun, target_dims = 'time', output_dims = 'time')$output1
Season(1 : 59, posdim = 1, monini = 1, moninf = 3, monsup = 5)
res_season <- Apply(list(a), fun = Season, posdim = 4, monini = 1, moninf = 3, 
                    monsup = 5, target_dims = 'time',
                    output_dims = 'time')$output1
res_season <- Apply(list(a), fun = Season, posdim = 2, monini = 1, moninf = 3, 
                    monsup = 5, target_dims = 'time',
                    output_dims = 'time')$output1

res_season <- Apply(list(a), fun = Season, posdim = 4, monini = 1, moninf = 3, 
                    monsup = 5, target_dims = c('dat', 'time'),
                    output_dims = 'time')$output1
res_season <- Apply(list(a), fun = Season, posdim = 2, monini = 1, moninf = 3, 
                    monsup = 5, target_dims = c('dat', 'time'),
                    output_dims = 'time')$output1
res_season <- Apply(list(a), fun = Season, posdim = 2, monini = 1, moninf = 3, 
                    monsup = 5, margins = c('lat', 'lon'))$output1
res_season <- Apply(list(a), fun = Season, posdim = 2, monini = 1, moninf = 3,
                    monsup = 5, margins = c('dat'))$output1


res9 <- Compute(wf,
               chunks = list(ensemble = 2, sdate=2),
               threads_load = 1,
               threads_compute = 4,
               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 = 'startR/0.1.2-foss-2018b-R-3.5.0',
                              cores_per_job = 2,
                              job_wallclock = '00:10:00',_
                              max_jobs = 4,
                              #extra_queue_params = list('#SBATCH --mem-per-cpu=3000'),
                              extra_queue_params = list('#SBATCH --qos=bsc_es'),
                              bidirectional = FALSE,
                              polling_period = 10),
               ecflow_suite_dir = '/esarchive/scratch/nperez/ecFlow/') #,
               #wait = TRUE)

                  
