#--------------------------------------------------------
# Author: Carlos Delgado, An-Chi Ho
# Date: 24th March 2021
#--------------------------------------------------------
# Example on how to load monthly data from decadal prediction systems that are 
# initialised in different dates:
# - HadGEM3 (initialised in November)
# - IPSL-CM6A-LR (initialised in January)
# - NorCPM1 (initialised in October)
# Therefore, we use multiStart() to insert NAs at missing months manually. 
#
# To see more details of multiStart() usage, check ex2_11.
#--------------------------------------------------------
library(startR)

variable <- 'tas'
sdates <- 2000:2004

## Paths
path_hadgem3 <- paste0('/esarchive/exp/CMIP6/dcppA-hindcast/hadgem3-gc31-mm/cmip6-dcppA-hindcast_i1p1/DCPP/',
                       'MOHC/HadGEM3-GC31-MM/dcppA-hindcast/r1i1p1f2/Amon/$var$/gr/v20200316/',
                       '$var$_*_s$sdate$_r1i1p1f2_gr_$fyear$.nc')
path_ipsl <- paste0('/esarchive/exp/CMIP6/dcppA-hindcast/ipsl-cm6a-lr/cmip6-dcppA-hindcast_i1p1/DCPP/',
                    'IPSL/IPSL-CM6A-LR/dcppA-hindcast/r1i1p1f1/Amon/$var$/gr/v20200504/',
                    '$var$_Amon_IPSL-CM6A-LR_dcppA-hindcast_s$sdate$-r1i1p1f1_gr_$fyear$.nc')
path_norcpm1 <- paste0('/esarchive/exp/CMIP6/dcppA-hindcast/norcpm1/cmip6-dcppA-hindcast_i1p1/DCPP/',
                       'NCC/NorCPM1/dcppA-hindcast/r1i1p1f1/Amon/$var$/gn/v20200320/',
                       '$var$_Amon_NorCPM1_dcppA-hindcast_s$sdate$-r1i1p1f1_gn_$fyear$.nc')

## Forecast months: six indices (from October to March)
time_hadgem3 <- indices(c(NA, 1, 2, 3, 4, 5))
time_ipsl <- indices(c(NA, NA, NA, 1, 2, 3))
time_norcpm1 <- indices(c(1, 2, 3, 4, 5, 6))

data <- multiStart(dat = list(list(name = 'hadgem3', path = path_hadgem3),
                              list(name = 'ipsl', path = path_ipsl),
                              list(name = 'norcpm1', path = path_norcpm1)),
                   var = variable,
                   sdate = paste0(sdates),
                   fyear = indices(1), ## there is one file per sdate
                   fyear_depends = 'sdate',
                   time = list(list(name = 'hadgem3', time = time_hadgem3),
                               list(name = 'ipsl', time = time_ipsl),
                               list(name = 'norcpm1', time = time_norcpm1)),
                   lat = values(list(-30, 30)),
                   lon = values(list(0, 50)),
                   synonims = list(lon = c('lon', 'longitude'), lat = c('lat', 'latitude')),
                   transform = CDORemapper,
                   transform_params = list(grid = 'r180x90',
                                           method = 'conservative',
                                           crop = c(0, 50, -30, 30)),
                   transform_vars = c('lat', 'lon'),
                   return_vars = list(lat = NULL, lon = NULL, time = 'sdate'),
                   retrieve = FALSE)


# Notice that though multiStart() returns one object 'data', it actually contains 
# three startR calls.
length(data)
# 3
names(data)
#[1] "hadgem3" "ipsl"    "norcpm1"

# Therefore, the operation should consider three data arrays 
# rather than one.
fun <- function(dat1, dat2, dat3) {
  # dat: [dat = 1]
  mean(c(dat1, dat2, dat3), na.rm = TRUE)
}

# As explained above, the number of data arrays is three, so the target_dims is a 
# list of three. The names of dataset (i.e., dat1, dat2, dat3) can be random.
step <- Step(fun = fun,
             target_dims = list(dat1 = 'dat', dat2 = 'dat', dat3 = 'dat'),
             output_dims = NULL)

# The inputs is still 'data', which contains three calls implicitly.
wf <- AddStep(inputs = data, step_fun = step)


# Option 1: run locally
output <- Compute(workflow = wf,
                  chunks = list(sdate = 5))

# Option 2: run on Nord3
#-------------------user-defined---------------------
  queue_host <- 'nord3'
  temp_dir <- '/gpfs/scratch/bsc32/bsc32<id>/startR_hpc/'
  ecflow_suite_dir <- '/home/Earth/<use_id>/startR_local/'
#----------------------------------------------------
output <- Compute(workflow = wf,
                  chunks = list(sdate = 5),
                  threads_load = 1,
                  threads_compute = 5,
                  cluster = list(queue_host = queue_host,
                                 queue_type = 'lsf',
                                 temp_dir = temp_dir,
                                 cores_per_job = 1,
                                 job_wallclock = '01:00',
                                 max_jobs = 5,
                                 extra_queue_params = list('#BSUB -q bsc_es -M 3000'), 
                                 bidirectional = FALSE,
                                 polling_period = 10),
                  ecflow_suite_dir = ecflow_suite_dir,
                  wait = TRUE)


# Results
dim(output$output1)
# time   var sdate fyear   lat   lon 
#    6     1     5     1    30    26 

summary(output$output1)
#   Min. 1st Qu.  Median    Mean 3rd Qu.    Max. 
#  279.4   293.4   296.6   295.5   298.7   304.1
