#----------------------------------------------------------------------------
# Author: An-Chi Ho
# Date: 20th May 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)

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

data <- Start(dataset = path,
              var = 'tos',
              sdate = c('1960', '1961'),
              aux = 'all',
              aux_depends = 'sdate',
              j = 'all', 
              i = 'all', 
              time = indices(1:12),
              member = 'r1i1p1f1',
              return_vars = list(j = NULL, i = NULL, 
                                 latitude = NULL, longitude = NULL),
              retrieve = F)

func_regrid <- function(data) {
  lons <- attr(data, 'Variables')$common$longitude
  lats <- attr(data, 'Variables')$common$latitude
  data <- s2dv::CDORemap(data, lons, lats, 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))
}

#NOTE: The data transposes if target_dims are only 'j' and 'i'. 
#      If only 'j' and 'i', output_dims will be 'lat', 'lon'.
step <- Step(fun = func_regrid,
             target_dims = list(data = c('var', 'aux', 'j', 'i')),
             output_dims = list(data = c('var', 'aux','lon', 'lat'), 
                                lats = 'lat', lons = 'lon'),
             use_attributes = list(data = "Variables"))
wf <- AddStep(data, step)

res <- Compute(workflow = wf$data,
               chunks = list(sdate = 2, time = 2))

names(res)
#[1] "data" "lats" "lons"
dim(res$data)
#    lat     lon dataset     var   sdate     aux    time  member 
#    180     360       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 

s2dv::PlotEquiMap(drop(res$data)[ , , 1, 1],
                  lon = drop(res$lons)[, 1, 1],
                  lat = drop(res$lats)[, 1, 1], fileout = '~/irre.png')

s2dv::PlotEquiMap(t(res$data[1, 1, 50:90, 30:50, 1, 1, 1, 1]),
                  lon = res$lons[50:90, 1, 1, 1, 1],
                  lat = res$lats[30:50, 1, 1, 1, 1], 
                  bar_limits = c(-2, 32), fileout = '~/irre2.png')
