
# -------------------------- CARLOS CODE

library(startR)
path = '/esarchive/exp/ecearth/a1ua/cmorfiles/DCPP/EC-Earth-Consortium/EC-Earth3/dcppA-hindcast/$member$/Omon/$var$/gr/v20190713/$var$_Omon_EC-Earth3_dcppA-hindcast_s$sdate$-$member$_gr_$fyear$.nc'

  ## Regions for AMV (Doblas-Reyes et al., 2013)

  regions = NULL
  regions$reg1 = c(0, 60, 0, 40)

 data <- Start(dataset = path,
                     var = 'tos',
                     sdate = paste0(1990:1992),
                     time = c(1:5),
                     lat = values(list(0, 20)),
                     lon = values(list(0, 40)),
                     fyear = 'all',
                     member = 'all',
                     fyear_depends = 'sdate',
                     time_across = 'fyear',
                     merge_across_dims = TRUE,
                     synonims = list(time = c('fmonth','time'),
                                     lon = c('lon', 'longitude'),
                                     lat = c('lat', 'latitude')),
                     return_vars = list(lat = 'dataset', lon = 'dataset'),
                     num_procs = 4, retrieve = FALSE)

data_lats <- as.vector(attributes(data)$Variables$dat1$lat)
data_lons <-  as.vector(attributes(data)$Variables$dat1$lon)


fun_amv_ipo <- function(data, data_lats, data_lons, index, regions) {
  
  data = s2dv::Season(data = data, time_dim = 'time',
                      monini = 11, moninf = 1, monsup = 12,
                      method = mean, na.rm = TRUE)
  names(dim(data))[which(names(dim(data))=='fmonth')] = 'fyear'
  mean_1 = ClimProjDiags::WeightedMean(data = data, lon = data_lons,
               lat = data_lats, region = regions$reg1,
               londim = which(names(dim(data))=='lon'),
               latdim = which(names(dim(data))=='lat'))
return(mean_1)
}

step <- Step(fun = fun_amv_ipo,
            #use_libraries = c('startR','s2dv','multiApply','ClimProjDiags'),
            target_dims = c('dataset','var','sdate','time','lat','lon'),
            output_dims = c('dataset','var','sdate','fyear'))

wf <- AddStep(inputs = data,
             data_lats = data_lats,
             data_lons = data_lons,
             index = 'AMV',
             regions = regions,
             step = step)

res <- Compute(workflow = wf,
               chunks = list(member = 10),
               threads_load = 2,
               threads_compute = 4,
               cluster = list(queue_host = 'nord3',
                              queue_type = 'lsf',
                        temp_dir = '/gpfs/scratch/bsc32/bsc32339/startR_hpc/',
                              cores_per_job = 2,
                              job_wallclock = '01:00',
                              max_jobs = 10,
                              extra_queue_params = list('#BSUB -q bsc_es'),
                              bidirectional = FALSE,
                              polling_period = 10),
               ecflow_suite_dir = '/home/Earth/nperez/startR_local/',
               wait = TRUE)
res_WS <- Compute(workflow = wf, chunks = list(member = 10))
range(res_WS)
#[1] 299.8463 301.7826


res <- Compute(workflow = wf,
               chunks = list(member = 10),
               threads_load = 2,
               threads_compute = 4,
               cluster = list(queue_host = 'cte-power',
                              queue_type = 'slurm',
                        temp_dir = '/gpfs/scratch/bsc32/bsc32339/startR_hpc/',
                              cores_per_job = 2,
                              job_wallclock = '01:00:00',
                              max_jobs = 10,
                              r_module = 'R/3.6.1-foss-2018b', #startR 1.0.0
                              bidirectional = FALSE,
                              polling_period = 10),
               ecflow_suite_dir = '/home/Earth/nperez/startR_local/',
               wait = TRUE)



