#----------------------------------------------------------------------------
# Author: An-Chi Ho
# Date: 8th Oct 2021
#
# This script shows how to load irregular grid data by Start(), then regrid it
# by s2dv::CDORemap in the workflow. It is a solution before Start() can deal
# with irregular regridding directly.
#----------------------------------------------------------------------------

library(startR)
library(s2dv)

path <- paste0('/esarchive/exp/CMIP6/dcppA-hindcast/CMCC-CM2-SR5/',
               'DCPP/CMCC/CMCC-CM2-SR5/dcppA-hindcast/$member$/Omon/$var$/gn/v20210719/',
               '$var$_*_s$sdate$-$member$_gn_$aux$.nc')

data <- Start(dataset = path,
              var = 'tos',
              sdate = c('1960', '1961'),
              aux = 'all',
              aux_depends = 'sdate',
              i = indices(2:291), # remove two indices to avoid white strips
              j = indices(2:361), # remove two indices to avoid white strips
              time = indices(1:12),
              member = 'r1i1p1f1',
              return_vars = list(latitude = NULL, longitude = NULL),
              retrieve = F)

attr(data, 'Dimensions')
#dataset     var   sdate     aux       i       j    time  member 
#      1       1       2       1     290     360      12       1 
dim(attr(data, 'Variables')$common$longitude)
#  j   i 
#362 292 
dim(attr(data, 'Variables')$common$latitude)
#  j   i 
#362 292 

func_regrid <- function(data) {
  lons <- attr(data, 'Variables')$common$longitude
  lats <- attr(data, 'Variables')$common$latitude
  data <- s2dv::CDORemap(data, lons[2:361, 2:291], lats[2:361, 2:291],
                         grid = 'r360x180', method = 'bil', crop = FALSE)
  lons_reg <- data[['lons']]
  lats_reg <- data[['lats']]
  return(list(data = data[[1]], lats = lats_reg, lons = lons_reg))
}

step <- Step(fun = func_regrid,
             target_dims = list(data = c('i', 'j')),
             output_dims = list(data = c('lon', 'lat'), 
                                lats = 'lat', lons = 'lon'),
             use_attributes = list(data = "Variables"))
wf <- AddStep(data, step)



#========= OPTION 1: Compute locally =============================
res <- Compute(workflow = wf$data,
               chunks = list(sdate = 2, time = 2))


#========= OPTION 2: Compute ON NORD4 with AUTOSUBMIT ============
#-----------modify according to your personal info-----------
hpc_user <- "bsc032xxx"
expid <- "xxxx" # autosubmit exp id; can be NULL
autosubmit_suite_dir <- "/home/Earth/<user_id>/startR_local_autosubmit/"
#------------------------------------------------------------
res <- Compute(workflow = wf$data, 
               chunks = list(sdate = 2, time = 2),
               threads_compute = 4,
               threads_load = 2,
               cluster = list(
                 queue_host = 'nord4',
                 expid = expid,
                 hpc_user = hpc_user,
                 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 = 4,
                 job_wallclock = '01:00:00',
                 max_jobs = 4,
                 special_setup = 'nord4'
               ),
               workflow_manager = 'autosubmit',
               autosubmit_suite_dir = autosubmit_suite_dir,
               autosubmit_server = NULL,
               wait = TRUE)



names(res)
#[1] "data" "lats" "lons"
dim(res$data)
#    lon     lat dataset     var   sdate     aux    time  member 
#    360     180       1       1       2       1      12       1 
dim(res$lons)
#    lon dataset     var   sdate     aux    time  member 
#    360       1       1       2       1      12       1 
dim(res$lats)
#    lat dataset     var   sdate     aux    time  member 
#    180       1       1       2       1      12       1 

PlotEquiMap(drop(res$data)[ , , 1, 1],
            lon = drop(res$lons)[, 1, 1],
            lat = drop(res$lats)[, 1, 1])

# Plot Layout for sdate = 1 all the time steps
var <- Reorder(drop(res$data)[, , 1, ], c(3, 1, 2))
PlotLayout(PlotEquiMap, c('lon', 'lat'), var = var, 
           lon = drop(res$lons)[, 1, 1], lat = drop(res$lats)[, 1, 1])

