# ------------------------------------------------------------------------------
# Downscaling precipitation using RainFARM 
# ------------------------------------------------------------------------------
# This usecase demonstrates that the start date dimension is used as chunking 
# dimension, but the chunk number is used to know the start date value of each
# chunk.
# The first part of the function performs downscaling method, which requires 
# longitude and latitude dimensions, so these two dimensions must be the target
# dimensions in the workflow.
# After that, the results are saved as netCDF file following esarchive convention.
# We need start date value here to decide the file name.
# As you can see, the sdate dimension is not required for the computation, so it 
# is not necessary to be the target dimension. We can just use 'chunk_indices' to
# get the chunk number therefore get the corresponding start date value for the 
# file name.
# ------------------------------------------------------------------------------
# Note 1: The data could be first transformed with QuantileMapping from CSTools
# Note 2: Extra parameters could be used to downscale the data: weights, slope...
# See more information in: 
# https://cran.r-project.org/web/packages/CSTools/vignettes/RainFARM_vignette.html
# ------------------------------------------------------------------------------
# MOST IMPORTANT NOTE:
# startR aims to return a result that fits in your local memory. This aim is 
# the oposite of downscaling, which increases the output size. Therefore, this 
# example saves the data to NetCDF files and provides the mean which has a 
# reduced size.
# Warning!!!!!! Use this example with caution to avoid saving not desired data.
#-------------------------------------------------------------------------------

# Load libraries and functions:
library(startR)
library(CSTools)
library(ncdf4)  
library(s2dv)

# Define the data:
sdates <- paste0(1993:1996, '1101')  # starting dates
path <- "/esarchive/exp/ecmwf/system5c3s/daily/$var$_s0-24h/$var$_$sdate$.nc"
data <- Start(dataset = path,
              var = 'prlr',
              sdate = sdates,
              member = 1:3,
              longitude = values(list(-10, 29)),
              longitude_reorder = CircularSort(-180, 180),
              latitude = values(list(18, 57)),
              time = 'all', 
              return_vars = list(latitude = 'dataset', longitude = 'dataset',
                                 time = NULL),
              synonims = list(latitude = c('lat', 'latitude'),
                              longitude = c('lon', 'longitude'),
                              member = c('member', 'ensemble')),
              retrieve = FALSE)

# Define the function:
Chunk_RF <- function(x, nf, destination, startdates, nchunks = chunk_indices) {
  lon <- as.numeric(attributes(x)$Variables$dat1$longitude)
  lat <- as.numeric(attributes(x)$Variables$dat1$latitude)
  down_data <- RainFARM(x, lon = lon, lat = lat, drop_realization_dim = TRUE,
                        nf, lon_dim = 'longitude', lat_dim = 'latitude', time_dim = 'time')
  # detect the dates of forecast time for different start dates
  time <- attributes(x)$Variables$common$time 
  dates <- lapply(startdates, function(x) {seq(as.Date(x, format = "%Y-%m-%d"), 
                                                       x + length(time) - 1, 'day')})
  dimname <- names(dim(down_data$data)) 
  var_dims <- list(ncdim_def(name = 'lon', units = 'degrees',
                             vals = as.vector(down_data$lon), longname = 'longitude'),
                   ncdim_def(name = 'lat', units = 'degrees',
                             vals = as.vector(down_data$lat), longname = 'longitude'),
                   ncdim_def(name = 'ensemble', units = 'adim',
                             vals = 1 : dim(down_data$data)[which(dimname == 'member')],
                             longname = 'ensemble', create_dimvar = TRUE))
  metadata_var <- list(units = 'm s-1', 
                       cdo_grid_name = paste0('r', length(lon), 'x', length(lat)), 
                       projection = 'none')
  # startdates and dates depends on the chunk:
  CSTools:::.saveExp(down_data$data, startdates =  startdates[nchunks['sdate']], 
                     dates = dates[[nchunks['sdate']]], defined_dims = var_dims, 
                     varname = 'prlr', metadata_var = metadata_var, 
                     destination = destination)
  down_data_mean <- s2dv::MeanDims(down_data$data, c('member', 'longitude', 'latitude'))
  return(down_data_mean)
}

step <- Step(Chunk_RF,
             target_dims = c('member', 'longitude', 'latitude', 'time'),
             output_dims = 'time',
             use_libraries = c('CSTools', 'ncdf4'),
             use_attributes = list(data = "Variables"))

workflow <- AddStep(list(data = data), step, nf = 4,
                    destination = "./test_RF_start/",
                    startdates = as.Date(sdates, format = "%Y%m%d"))
  

#========= OPTION 1: Compute locally ============
res <- Compute(workflow,
               chunks = list(sdate = 4),
               threads_load = 2,
               threads_compute = 4)


#========= OPTION 2: Compute ON NORD4 ============

#-----------modify according to your personal info---------
  queue_host = 'nord4'   # your own host name for nord4
  temp_dir =  '/gpfs/scratch/bsc32/bsc032339/startR_hpc/'
  ecflow_suite_dir = '/home/Earth/nperez/startR_local/'  # your own local directory
#------------------------------------------------------------

res <- Compute(workflow,
               chunks = list(sdate = 4),
               threads_load = 1,
               threads_compute = 1,
               cluster = list(queue_host = queue_host,
                              queue_type = 'slurm',
                              cores_per_job = 16,
                              r_module = 'R/4.1.2-foss-2019b',
                              temp_dir = temp_dir,
                              polling_period = 10,
                              job_wallclock = '01:00:00',
                              max_jobs = 4,
                              bidirectional = FALSE,
                              special_setup = 'nord4',
                              extra_queue_params = list('#SBATCH -A bsc32', '#SBATCH -q bsc_es')),
               ecflow_suite_dir = ecflow_suite_dir,
               wait = TRUE)

# Visualize the temporal evolution of the result simultaneously for all sdates:
matplot(1:215, res$output1[, 1, 1, ], type = "l")
