# -----------------------------------------------------------------
# Function working on time dimension e.g.: Season
# ------------------------------------------------------------------
rm(list = ls())
library(startR)
packageVersion("startR")

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 = FALSE)

library(multiApply)
fun_spring <- function(x) {
    source("/esarchive/scratch/nperez/Season_v2.R")
    y <- Season_v2(x,  monini = 1, moninf = 3, monsup = 5)
    return(y)
}

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

wf1 <- AddStep(data, step1)

# Locally
res1 <- Compute(wf1,
            chunks = list(ensemble = 2,
                          sdate = 2))

dim(res1$output1)
str(res1$output1)
summary(res1$output1)
# -----------------------------------------------------------------
#> dim(res1$output1)
#     time       dat       var     sdate  ensemble  latitude longitude 
#        1         1         1         2        20       640        40 
#> str(res1$output1)
# num [1, 1, 1, 1:2, 1:20, 1:640, 1:40] 260 257 257 260 257 ...
#> summary(res1$output1)
#   Min. 1st Qu.  Median    Mean 3rd Qu.    Max. 
#  220.5   273.3   282.7   279.9   295.6   306.8 
# ----------------------------------------------------------------
res1_P <- Compute(wf1,
                  chunks = list(ensemble = 20,
                                sdate = 2),
                  threads_load = 2,
                  threads_compute = 4,
                  cluster = list(queue_host = 'cte-power',
                                 queue_type = 'slurm',
                                 cores_per_job = 2,
                                 temp_dir = '/gpfs/scratch/bsc32/bsc32339/',
                                 r_module = 'R/3.5.0-foss-2018b',
                     extra_queue_params = list('#SBATCH --mem-per-cpu=3000'),
                                 polling_period = 10,
                                 job_wallclock = '01:00:00',
                                 max_jobs = 40,
                                 bidirectional = FALSE),
               ecflow_suite_dir = '/esarchive/scratch/nperez/',
               wait = TRUE)

#-----------------------------------------------------------------
# Function using attributes of the data e.g.: latitudinal correction
# ------------------------------------------------------------------
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 = FALSE)
funp <- function(x) {
          lat = attributes(x)$Variables$dat1$latitude
          weight = sqrt(cos(lat * pi / 180))
          corrected = Apply(list(x), target_dims = "latitude",
                            fun = function(x) {x * weight})
        }


step2 <- Step(fun = funp,
             target_dims = 'latitude',
             output_dims = 'latitude',
             use_attributes = list(data = "Variables"))

wf2 <- AddStep(data, step2)
res2 <- Compute(workflow = wf2,
            chunks = list(ensemble = 2,
                          sdate = 2))

dim(res2$output1)
head(res2$output1)
summary(res2$output1)
# ------------------------------------------------------------------
# Output:
#> dim(res2$output1)
# latitude       dat       var     sdate  ensemble      time longitude
#      640         1         1         2        20         7        40
#> head(res2$output1)
#[1] 15.22683 23.09543 28.94978 33.82982 38.11695 41.99680
#> summary(res2$output1)
#   Min. 1st Qu.  Median    Mean 3rd Qu.    Max.
#  13.19  169.60  237.60  217.90  284.10  305.10
# -----------------------------------------------------------------

res2_P <- Compute(wf2,
                  chunks = list(ensemble = 20,
                                sdate = 2),
                  threads_load = 2,
                  threads_compute = 4,
                  cluster = list(queue_host = 'cte-power',
                                 queue_type = 'slurm',
                                 cores_per_job = 2,
                                 temp_dir = '/gpfs/scratch/bsc32/bsc32339/',
                                 r_module = 'R/3.5.0-foss-2018b',
                     extra_queue_params = list('#SBATCH --mem-per-cpu=3000'),
                                 polling_period = 10,
                                 job_wallclock = '01:00:00',
                                 max_jobs = 40,
                                 bidirectional = FALSE),
               ecflow_suite_dir = '/esarchive/scratch/nperez/',
               wait = TRUE)


# --------------------------------------------------------------
# Function doing regridding CDO (e.g.: s2dverification::CDORemap)
#---------------------------------------------------------------
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 : 2),
              time = 'all',
              latitude = 'all', #indices(1:40),
              longitude = 'all', #indices(1 : 40),
              return_vars = list(latitude = 'dat', longitude = 'dat', time = 'sdate'),
              retrieve = FALSE)



fun_deb2 <- function(x) {
             lons_data = as.vector(attr(x, 'Variables')$dat1$longitude)
             lats_data = as.vector(attr(x, 'Variables')$dat1$latitude)
             resgrid = "r360x180"   # prlr
             r <- s2dverification::CDORemap(x, lons_data, lats_data, resgrid,
                                      'bil', crop = FALSE, force_remap = TRUE)[[1]]
             return(r)
}

step3 <- Step(fun = fun_deb2,
             target_dims = c('latitude','longitude'),
             output_dims = c('latitude', 'longitude'),
             use_attributes = list(data = "Variables"))
wf3 <- AddStep(data, step3)
res3 <- Compute(workflow = wf3,
            chunks = list(ensemble = 2,
                          sdate = 2))
dim(res3$output1)
head(res3$output1)
summary(res3$output1)

# --------------------------------------------------------------------
# Output:
#> dim(res3$output1)
# latitude longitude       dat       var     sdate  ensemble      time
#      180       360         1         1         2         2         7
#> head(res3$output1)
#[1] 245.2156 244.6700 245.1513 244.7413 244.4198 245.0041
#> summary(res3$output1)
#   Min. 1st Qu.  Median    Mean 3rd Qu.    Max.
#  206.3   269.2   283.3   278.8   296.4   314.2
# --------------------------------------------------------------------

res3_P <- Compute(wf3,
                  chunks = list(ensemble = 2,
                                sdate = 2),
                  threads_load = 2,
                  threads_compute = 4,
                  cluster = list(queue_host = 'cte-power',
                                 queue_type = 'slurm',
                                 cores_per_job = 1,
                                 temp_dir = '/gpfs/scratch/bsc32/bsc32339/',
                                 r_module = 'R/3.5.0-foss-2018b',
                                 CDO_module = 'CDO/1.9.5-foss-2018b',
                     extra_queue_params = list('#SBATCH --mem-per-cpu=3000'),
                                 polling_period = 100,
                                 job_wallclock = '01:00:00',
                                 max_jobs = 4,
                                 bidirectional = FALSE),
               ecflow_suite_dir = '/esarchive/scratch/nperez/',
               wait = TRUE)


# --------------------------------------------------------------
# Two functions (e.g.: s2dverification::CDORemap and Season)
#---------------------------------------------------------------

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 : 2),
              time = 'all',
              latitude = 'all', #indices(1:40),
              longitude = 'all', #indices(1 : 40),
              return_vars = list(latitude = 'dat', longitude = 'dat', time = 'sdate'),
              retrieve = FALSE)

fun_deb3 <- function(x) {
             source("/esarchive/scratch/nperez/Season_v2.R")
             lons_data = as.vector(attr(x, 'Variables')$dat1$longitude)
             lats_data = as.vector(attr(x, 'Variables')$dat1$latitude)
             resgrid = "r360x180"   # prlr
             y = Season_v2(x, posdim = 'time', monini = 1, moninf = 1, monsup = 3)
             r <- s2dverification::CDORemap(y, lons_data, lats_data, resgrid,
                                      'bil', crop = FALSE, force_remap = TRUE)[[1]]
             return(r)
}


step4 <- Step(fun = fun_deb3,
             target_dims = c('latitude','longitude', 'time'),
             output_dims = c('latitude', 'longitude', 'time'),
             use_attributes = list(data = "Variables"))
wf4 <- AddStep(data, step4)
res4 <- Compute(workflow = wf4,
            chunks = list(ensemble = 2,
                          sdate = 2))
dim(res4$output1)
head(res4$output1)
summary(res4$output1)

# ------------------------------------------------------------------
# Output:
#> dim(res4$output1)
# latitude longitude      time       dat       var     sdate  ensemble
#      180       360         1         1         1         2         2
#> head(res4$output1)
#[1] 237.1389 237.2601 238.0882 238.0312 237.7883 238.4835
#> summary(res4$output1)
#   Min. 1st Qu.  Median    Mean 3rd Qu.    Max.
#  227.3   259.6   280.8   277.1   296.2   306.7
# ------------------------------------------------------------------
res4_P <- Compute(wf4,
                  chunks = list(ensemble = 2,
                                sdate = 2),
                  threads_load = 2,
                  threads_compute = 4,
                  cluster = list(queue_host = 'cte-power',
                                 queue_type = 'slurm',
                                 cores_per_job = 1,
                                 temp_dir = '/gpfs/scratch/bsc32/bsc32339/',
                                 r_module = 'R/3.5.0-foss-2018b',
                                 CDO_module = 'CDO/1.9.5-foss-2018b',
                     extra_queue_params = list('#SBATCH --mem-per-cpu=3000'),
                                 polling_period = 10,
                                 job_wallclock = '01:00:00',
                                 max_jobs = 6,
                                 bidirectional = FALSE),
               ecflow_suite_dir = '/esarchive/scratch/nperez/',
               wait = TRUE)

