#Clear workspace
rm(list = ls())

.libPaths(new = '/gpfs/projects/bsc32/share/R_libs/3.5') 
library(maps)
library(boot)
library(RColorBrewer)
library(easyVerification)
library(SpecsVerification)
library(s2dverification)
library(abind)
library(multiApply)
library(ncdf4)
library(startR)
library(plyr)
source('/esarchive/scratch/Earth/rfernand/dcpp/R_WeightedMean.R')

#-----------------

vari = 'tos'
mod = 'Omon'
grid = 'gr'

#Load CMIP6 DCPP
repos <- paste0('/esarchive/exp/ecearth/a1ua/original_files/cmorfiles_corrupt/DCPP/EC-Earth-Consortium/EC-Earth3/dcppA-hindcast/$memb$/',mod,'/$var$/gr/v20190713/$var$_',mod,'_EC-Earth3_dcppA-hindcast_s$sdate$-$memb$_',grid,'_$chunk$.nc',sep='')


#=== 1. subset data & retrieve = TRUE ===
data_cmip6_1 <- Start(dat = repos,
              var = vari,
              memb = c('r1i1p1f1'),
              sdate = '2014', #paste(seq(1960,2015)),
              time = 'all', #c(1:62),
              chunk = 'all',
              lat = 'all',
              lon = 'all',
              chunk_depends = 'sdate',
              time_across = 'chunk',
              merge_across_dims = TRUE,
              return_vars = list(lat = 'dat', lon = 'dat'),
              retrieve = T)

#Use Season() directly
    y <- Season_v1(data_cmip6_1,  posdim = 5, monini = 11, moninf = 1, monsup = 12)
    z <- Season_v2(data_cmip6_1, posdim = 'time', monini = 11, moninf = 1, monsup = 12)

#>>> no error


#=== 2. subset data & retrieve = FALSE & run on ws===
data_cmip6_1 <- Start(dat = repos,
                      var = vari,
                      memb = c('r1i1p1f1'),
                      sdate = '2014',#c('2012','2013','2014','2015'), #paste(seq(1960,2015)),
                      time = 'all', #c(1:62),
                      chunk = 'all',
                      lat = 'all',
                      lon = 'all',
                      chunk_depends = 'sdate',
                      time_across = 'chunk',
                      merge_across_dims = TRUE,
                      return_vars = list(lat = 'dat', lon = 'dat'),
                      retrieve = F)

fun_an <- function(x) {
    #print(dim(x))
source("/esarchive/scratch/nperez/Season_v2.R")
    y <- Season_v2(x, posdim = 'time', monini = 11, moninf = 1, monsup = 12)
#    y <- s2dverification::Season(x, posdim = 1, monini = 11, moninf = 1, monsup = 12)
#    y <- .Season(x, monini = 11, moninf = 1, monsup = 12)  #also works
    return(y)
}

step1 <- Step(fun = fun_an,
             target_dims = c('time'),
             output_dims = c('time'))#,

wf <- AddStep(data_cmip6_1, step1)

print('Start Compute')

#run locally
#  res <- Compute(wf,
#               chunks = list(sdate = 1))
##                 chunks = list(lat = 2,
##                               lon = 2))

#Compute() on fatnodes
   res <- Compute(wf,
                  chunks = list(latitude = 2,
                                longitude = 2), $output1
                  threads_load = 2,
                  threads_compute = 4,
                  cluster = list(queue_host = 'fat',
                                 queue_type = 'slurm',
                                 cores_per_job = 2,
                                 job_wallclock = '00:10:00',
                                 max_jobs = 4,
                                 bidirectional = TRUE
                  ),
                  ecflow_suite_dir = '/home/Earth/aho/startR_local/'
                   wait = TRUE
          )

saveRDS(res, file = 'test_collect_rfernand.Rds')
