# ------------------------------------------------------------------
# Function working on time dimension e.g.: Season
# ------------------------------------------------------------------
rm(list = ls())
library(startR)
library(tictoc)
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)

devtools::install_git('https://earth.bsc.es/gitlab/es/s2dverification', branch = 'devel-SeasonApply')

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
tic()
res <- Compute(wf,
            chunks = list(ensemble = 2,
                          sdate = 2))
toc()

dim(res$output1)
str(res$output1)
summary(res$output1)

# Output:

> dim(res$output1)
     time       dat       var     sdate  ensemble  latitude longitude 
        1         1         1         2        20       640        40 
> str(res$output1)
 num [1, 1, 1, 1:2, 1:20, 1:640, 1:40] 260 257 257 260 257 ...
> summary(res$output1)
   Min. 1st Qu.  Median    Mean 3rd Qu.    Max. 
  220.5   273.3   282.7   279.9   295.6   306.8 



# ------------------------------------------------------------------
# Function using attributes of the data e.g.: latitudinal correction
# ------------------------------------------------------------------
rm(list = ls())
library(multiApply)
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)
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})
        }


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

wf <- AddStep(data, step)
tic()
res <- Compute(workflow = wf,
            chunks = list(ensemble = 2,
                          sdate = 2))
toc()

dim(res$output1)
head(res$output1)
summary(res$output1)

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




# --------------------------------------------------------------
# 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 = TRUE)
              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)
}


step <- Step(fun = fun_deb2,
             target_dims = c('latitude','longitude'),
             output_dims = c('latitude', 'longitude'),
             use_attributes = list(data = "Variables"))
wf <- AddStep(data, step)
tic()
res <- Compute(workflow = wf,
            chunks = list(ensemble = 2,
                          sdate = 2))
toc()

dim(res$output1)
head(res$output1)
summary(res$output1)

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

