# Author: An-Chi Ho
# Date: 24th Aug. 2020
# ------------------------------------------------------------------
# This use case shows you how to apply the existing mask file on your data.
# Case 1 is a land-sea mask, and case 2 is a mask file containing several
# region's mask.
# If you need to create the mask file on your own, go to ex2_9_mask.R


#================= Case 1 ==================
# Goal: Apply a land-sea mask on tas data. Because the mask is tailored for 
#       system5c3s data, the latitude and longitude of the mask and data 
#       have matched already.

# A) Load mask file
mask_path <- '/esarchive/autosubmit/con_files/land_sea_mask_1x1_c3s.nc'

mask <- Start(dat = mask_path,
              var =  'LSM',
              lat = 'all', 
              lon = 'all', 
              return_vars = list(var_names = NULL,
                                 lat = NULL, lon = NULL), 
              var_var = 'var_names', 
              retrieve = F)


# B) Load data
data_path <- '/esarchive/exp/ecmwf/system5c3s/monthly_mean/$var$_f6h/$var$_$sdate$.nc'

data <- Start(dat = data_path,
              var = 'tas',
              sdate = paste0(1991:2010, '0101'),
              ensemble = 'all',
              time = indices(1),
              lat = 'all',
              lon = 'all',
              synonims = list(lat = c('lat', 'latitude'),
                              lon = c('lon', 'longitude')), 
              return_vars = list(lat = NULL,
                                 lon = NULL,
                                 time = 'sdate'),
              retrieve = FALSE)

# C) Build the workflow
reg_mask <- function(data, mask) {
  # data: [dat = 1, sdate = 20, ensemble = 25, time = 1]
  # mask: [dat = 1]
  if (mask == 1) {
    ind <- s2dv::MeanDims(data, c('dat', 'ensemble', 'time'))
  } else {
    ind <- array(rep(NA, dim(data)[2]), dim = c(sdate = dim(data)[2]))
  }
  return(ind)

}

step_mask <- Step(fun = reg_mask,
                  target_dims = list(data = c('dat', 'sdate', 'ensemble', 'time'),
                                     mask = c('dat')),
                  output_dim = 'sdate')

wf_mask <- AddStep(list(data, mask), step_mask)

# run locally
res <- Compute(workflow = wf_mask,
                chunks = list(lat = 2,
                              lon = 2))



# Submit to Nord4
#-------------------user-defined---------------------
  queue_host <- 'n4login1.bsc.es'
  temp_dir <- '/gpfs/scratch/bsc32/bsc032734/startR_hpc/'
  ecflow_suite_dir <- '/home/Earth/aho/startR_local/'
#----------------------------------------------------
  res <- Compute(wf_mask,
                 chunks = list(lat = 2,
                               lon = 2),#$output1
                 threads_load = 2,
                 threads_compute = 4,
                 cluster = list(queue_host = queue_host,
                                queue_type = 'slurm',
                                temp_dir = temp_dir,
                                cores_per_job = 2,
                                job_wallclock = '05:00',
                                max_jobs = 4,
                                bidirectional = FALSE,
                                polling_period = 10,
                                special_setup = 'nord4',
                                extra_queue_params = list('#SBATCH -A bsc32', '#SBATCH -q bsc_es')
                 ),
                 ecflow_suite_dir = ecflow_suite_dir,
                 wait = TRUE
                 ) 

# Submit to Nord3v2
#-------------------user-defined---------------------
  queue_host <- 'nord4'
  temp_dir <- '/gpfs/scratch/bsc32/bsc32734/startR_hpc/'
  ecflow_suite_dir <- '/home/Earth/aho/startR_local/'
#----------------------------------------------------
  res <- Compute(wf_mask,
                 chunks = list(lat = 2,
                               lon = 2),#$output1
                 threads_load = 2,
                 threads_compute = 4,
                 cluster = list(queue_host = queue_host,
                                queue_type = 'slurm',
                                temp_dir = temp_dir,
                                cores_per_job = 2,
                                job_wallclock = '05:00',
                                max_jobs = 4,
                                bidirectional = FALSE,
                                polling_period = 10
                 ),
                 ecflow_suite_dir = ecflow_suite_dir,
                 wait = TRUE
                 ) 
  
#-------------------nord4 with autosubmit---------------------
  res <- Compute(wf_mask,
                 chunks = list(lat = 2,
                               lon = 2),#$output1
                 threads_load = 2,
                 threads_compute = 4,
                 cluster = list(queue_host = queue_host,  
                                expid = "a96t",
                                hpc_user = "bsc979986",
                                r_module = "R/4.1.2-foss-2019b",
                                CDO_module = "CDO/1.9.8-foss-2019b",
                                autosubmit_module = 'autosubmit/4.1.13-foss-2023b-Python-3.11.5',
                                cores_per_job = 2,
                                job_wallclock = '05:00',
                                max_jobs = 4,
                                bidirectional = FALSE,
                                polling_period = 10,
                                special_setup = 'nord4',
                                extra_queue_params = list('#SBATCH -A bsc32', '#SBATCH -q bsc_es')),
                 workflow_manager = 'autosubmit',
                 autosubmit_suite_dir = "/home/Earth/tkariyat/startR_local_autosubmit/", # Change to yours
                 autosubmit_server = NULL,
                 wait = TRUE)

#-------- See the result -----------
str(res$output1)
# num [1:20, 1, 1:181, 1:360] 249 245 247 245 247 ...
summary(res$output1)
#   Min. 1st Qu.  Median    Mean 3rd Qu.    Max.    NA's 
#  236.8   273.9   286.5   282.4   297.1   304.4  441960 

# Check with plotting
library(s2dv)
s2dv::PlotLayout(PlotEquiMap, c('lat', 'lon'), res$output1[ , 1, , ], 
                 lon = seq(0, 359, 1), lat = seq(90, -90, -1))



#================= Case 2 ==================
# Goal: Apply analysis only on Baffin Bay and Greenland region.
# The mask of the two regions is from the same file, while we have to read them
# seperately because the dimensions have to be in line with the data. That is,
# the length of 'var' can only be 1.
# Another way is read data with repetitive "var = c('siconc', 'siconc')" and read mask
# as one object with "var = c('Baffin_Bay', 'Grnland')". The following workflow need
# to adjust accordingly.

# A) Load mask file
mask_path <- '/esarchive/autosubmit/con_files/mask.regions.HadISST.nc'

mask1 <- Start(dat = mask_path,
               var = 'Baffin_Bay',
               lat = 'all',
               lon = 'all',
               synonims = list(lat = c('lat', 'y'),
                               lon = c('lon', 'x')),
               return_vars = list(var_names = NULL,
                                  lat = NULL, lon = NULL),
               var_var = 'var_names',
               retrieve = F)

mask2 <- Start(dat = mask_path,
               var = 'Grnland',
               lat = 'all',
               lon = 'all',
               synonims = list(lat = c('lat', 'y'),
                              lon = c('lon', 'x')),
               return_vars = list(var_names = NULL,
                                  lat = NULL, lon = NULL),
               var_var = 'var_names',
               retrieve = F)


# B) Load data
# The original longitude range is [0, 360]. Reorder to [-180, 180] to be in line with masks.
data_path <- '/esarchive/exp/gfdl-cm4/cmip6-piControl/monthly_mean/sic-lonlat/sic_$sdate$.nc'
lons.min <- -179.5
lons.max <- 179.5
lats.min <- -90
lats.max <- 90

data <- Start(dat = data_path,
              var = 'siconc',
              sdate = paste0('0', 521:540, '0101'),
              ensemble = 'all',
              time = indices(1),
              lat = values(list(lats.min, lats.max)),
              lat_reorder = Sort(),
              lon = values(list(lons.min, lons.max)),
              lon_reorder = CircularSort(-180, 180),
              synonims = list(lat = c('lat', 'latitude'),
                              lon = c('lon', 'longitude')), 
              return_vars = list(lat = NULL,
                                 lon = NULL,
                                 time = 'sdate'),
              retrieve = FALSE)


# C) Build the workflow
reg_mask <- function(data, mask1, mask2) {
  # data: [var = 1, sdate = 20, ensemble = 1, time = 1]
  # mask1 and mask2: [var = 1]
#--------Remove when s2dv::Trend is updated -------
  source('/esarchive/scratch/aho/Trend_tmp.R')
#-----------------------------------------------

  if (mask1 == 1 | mask2 == 1) {
#------------- Change when s2dv::Trend is updated ----------------
    trend <- Trend(data, time_dim = 'sdate')$trend[2, , , ]
#    trend <- s2dv::Trend(data, time_dim = 'sdate')$trend[2, , , ]
#-----------------------------------------------------------------
    ind <- mean(data)
  } else {
    ind <-  NA
  }
  return(ind)
}

# The necessary target dim of data for the function is 'sdate' only, but those not
# in mask dimensions are needed to be included here to make chunking work.
step_mask <- Step(fun = reg_mask,
                  target_dims = list(data = c('var', 'sdate', 'ensemble', 'time'),
                                     mask1 = c('var'),
                                     mask2 = c('var')),
                  output_dim = NULL)

wf_mask <- AddStep(list(data, mask1, mask2), step_mask)

# Run locally
res <- Compute(workflow = wf_mask,
               chunks = list(lat = 2,
                             lon = 2))

# Submit to Nord4
#-------------------user-defined---------------------
  queue_host <- 'nord4'
  temp_dir <- '/gpfs/scratch/bsc32/bsc032734/startR_hpc/'
  ecflow_suite_dir <- '/home/Earth/aho/startR_local/'
#----------------------------------------------------
  res <- Compute(wf_mask,
                 chunks = list(lat = 2,
                               lon = 2),#$output1
                 threads_load = 2,
                 threads_compute = 4,
                 cluster = list(queue_host = queue_host,
                                queue_type = 'slurm',
                                temp_dir = temp_dir,
                                cores_per_job = 2,
                                job_wallclock = '05:00',
                                max_jobs = 4,
                                bidirectional = FALSE,
                                polling_period = 10,
                                special_setup = 'nord4',
                                extra_queue_params = list('#SBATCH -A bsc32', '#SBATCH -q bsc_es')
                 ),
                 ecflow_suite_dir = ecflow_suite_dir,
                 wait = TRUE
                 )
#-------------------user-defined---------------------
  queue_host = 'nord4'   #your own host name for nord4
  temp_dir =  '/gpfs/scratch/bsc32/bsc979986/startR_hpc/'
  ecflow_suite_dir = '/home/Earth/tkariyat/startR_local/'  #your own local directory
##----------------------------------------------------
  res <- Compute(wf_mask,
                 chunks = list(lat = 2,
                               lon = 2),#$output1
                 threads_load = 2,
                 threads_compute = 4,
                 cluster = list(queue_host = queue_host,  
                                expid = "a96t",
                                hpc_user = "bsc979986",
                                r_module = "R/4.1.2-foss-2019b",
                                CDO_module = "CDO/1.9.8-foss-2019b",
                                autosubmit_module = 'autosubmit/4.1.13-foss-2023b-Python-3.11.5',
                                cores_per_job = 2,
                                job_wallclock = '05:00',
                                max_jobs = 4,
                                bidirectional = FALSE,
                                polling_period = 10,
                                special_setup = 'nord4',
                                extra_queue_params = list('#SBATCH -A bsc32', '#SBATCH -q bsc_es')),
                 workflow_manager = 'autosubmit',
                 autosubmit_suite_dir = "/home/Earth/tkariyat/startR_local_autosubmit/", # Change to yours
                 autosubmit_server = NULL,
                 wait = TRUE)


#--------Check the result ------------
str(res$output1)
# num [1, 1:180, 1:360] NA NA NA NA NA NA NA NA NA NA ...
summary(res$output1)
#   Min. 1st Qu.  Median    Mean 3rd Qu.    Max.    NA's 
#   0.00   78.83   95.44   80.00   97.60   99.91   64093 

# Check with plotting. Only the two regions have values
s2dv::PlotEquiMap(res$output1[1, , ], lon = seq(-179.5, 179.5, 1), lat = seq(-89.5, 89.5, 1))

