# Apply an existing mask file on data
# Author: An-Chi Ho
# ------------------------------------------------------------------



# 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 = 'tos',
              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: [sdate = 20, ensemble = 25]
  # mask: [dat = 1]
  if (mask == 1) {
    trend <- s2dv::Trend(data, time_dim = 'sdate')$trend[2, , , ]
    ind <- mean(data)
  } else {
    ind <-  NA
  }
  return(ind)

}

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

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

res <- Compute(workflow = wf_mask,
                chunks = list(lat = 4,
                              lon = 4))



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

mask1 <- Start(dat = mask_path,
              var = 'Baffin_Bay', # c('Grnland', '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
data_path <- '/esarchive/exp/gfdl-cm4/cmip6-piControl/monthly_mean/sic-lonlat/sic_$sdate$.nc'

data <- Start(dat = data_path,
              var = 'siconc',
              sdate = paste0('0', 521:540, '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, mask1, mask2) {
  # data: [sdate = 20, ensemble = 25]
  # mask: [var = 2]
#--------Remove when s2dv::Trend is updated -------
  source('/esarchive/scratch/aho/Trend_tmp.R')
#-----------------------------------------------

  if (mask == 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)

}

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

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

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