# Author: Chihchung Chou, An-Chi Ho
# Date: 1st July 2021
# ------------------------------------------------------------------
# This use case uses experimental and the corresponding observational data to calculate 
# the temporal mean and spatial weighted mean. Notice that the spatial resolutions of the
# two datasets are different, but it still works because lat and lon are target dimensions.
# ------------------------------------------------------------------
library(startR)


# exp
repos_exp <- paste0('/esarchive/exp/ecearth/a1tr/cmorfiles/CMIP/EC-Earth-Consortium/',
                    'EC-Earth3/historical/r24i1p1f1/Amon/$var$/gr/v20190312/',
                    '$var$_Amon_EC-Earth3_historical_r24i1p1f1_gr_$sdate$01-$sdate$12.nc')

exp <- Start(dat = repos_exp,
             var = 'tas',
             sdate = as.character(c(2005:2008)),
             time = indices(1:3),
             lat = 'all',
             lon = 'all',
             synonims = list(lat = c('lat', 'latitude'),
                             lon = c('lon', 'longitude')),
             return_vars = list(lon = NULL,
                                lat = NULL,
                                time = 'sdate'),
             retrieve = FALSE)

lons_exp <- as.vector(attr(exp, 'Variables')$common$lon)
lats_exp <- as.vector(attr(exp, 'Variables')$common$lat)
dates_exp <- attr(exp, 'Variables')$common$time

attr(exp, 'Dimensions')
#  dat   var sdate  time   lat   lon 
#    1     1     4     3   256   512 
dates_exp
# [1] "2005-01-16 12:00:00 UTC" "2006-01-16 12:00:00 UTC"
# [3] "2007-01-16 12:00:00 UTC" "2008-01-16 12:00:00 UTC"
# [5] "2005-02-15 00:00:00 UTC" "2006-02-15 00:00:00 UTC"
# [7] "2007-02-15 00:00:00 UTC" "2008-02-15 12:00:00 UTC"
# [9] "2005-03-16 12:00:00 UTC" "2006-03-16 12:00:00 UTC"
#[11] "2007-03-16 12:00:00 UTC" "2008-03-16 12:00:00 UTC"


# obs
  path.obs <- '/esarchive/recon/ecmwf/era5/monthly_mean/$var$_f1h-r1440x721cds/$var$_$date$.nc'
  obs <- Start(dat = path.obs,
               var = 'tas',
               date = unique(format(dates_exp, '%Y%m')),
               time = values(dates_exp), 
               time_across = 'date',
               merge_across_dims = TRUE,
               split_multiselected_dims = TRUE,
               lat = 'all',
               lon = 'all',
               synonims = list(lon = c('lon', 'longitude'),
                               lat = c('lat', 'latitude')),
               return_vars = list(lon = NULL,
                                  lat = NULL,
                                  time = 'date'),
               retrieve = FALSE)

lons_obs <- as.vector(attr(obs, 'Variables')$common$lon)
lats_obs <- as.vector(attr(obs, 'Variables')$common$lat)
dates_obs <- attr(obs, 'Variables')$common$time


attr(obs, 'Dimensions')
#     data       var     sdate      time     lat      lon 
#        1         1         4         3     721     1440 
dates_obs
# [1] "2005-01-16 11:30:00 UTC" "2006-01-16 11:30:00 UTC"
# [3] "2007-01-16 11:30:00 UTC" "2008-01-16 11:30:00 UTC"
# [5] "2005-02-14 23:30:00 UTC" "2006-02-14 23:30:00 UTC"
# [7] "2007-02-14 23:30:00 UTC" "2008-02-15 11:30:00 UTC"
# [9] "2005-03-16 11:30:00 UTC" "2006-03-16 11:30:00 UTC"
#[11] "2007-03-16 11:30:00 UTC" "2008-03-16 11:30:00 UTC"


fun <- function(exp, obs, 
                lons_exp = lons_exp, lats_exp = lats_exp,
                lons_obs = lons_obs, lats_obs = lats_obs) {
  # exp
  e <- s2dv::MeanDims(drop(exp), 'time')
  sst.e <- ClimProjDiags::WeightedMean(e, lons_exp, lats_exp,
                                       londim = which(names(dim(e)) == 'lon'),
                                       latdim = which(names(dim(e)) == 'lat'))
  index.exp <- (sst.e - mean(sst.e))/sd(sst.e)
  
  # obs         
  o <- s2dv::MeanDims(drop(obs), 'time')
  sst.o <- ClimProjDiags::WeightedMean(o, lons_obs, lats_obs,
                                       londim = which(names(dim(o)) == 'lon'),
                                       latdim = which(names(dim(o)) == 'lat'))
  index.obs <- (sst.o - mean(sst.o))/sd(sst.o)

  # give dim name
  dim(index.exp) <- c(sdate = length(index.exp))
  dim(index.obs) <- c(sdate = length(index.obs))

  return(list(ind_exp = index.exp, ind_obs = index.obs))

}

# If ClimProjDiags::WeightedMean accepts two-dim input, 'sdate' can be margin dimension.
step <- Step(fun,
             target_dims = list(exp = c('sdate', 'time', 'lat', 'lon'),
                                obs = c('sdate', 'time', 'lat', 'lon')),
             output_dims = list(ind_exp = 'sdate', ind_obs = 'sdate'))

workflow <- AddStep(list(exp = exp, obs = obs), step,
                    lons_exp = lons_exp, lats_exp = lats_exp,
                    lons_obs = lons_obs, lats_obs = lats_obs)

res <- Compute(workflow$ind_exp,
               chunks = list(var = 1))

str(res)
#List of 2
# $ ind_exp: num [1:4, 1, 1] 1.195 0.422 -0.6 -1.017
# $ ind_obs: num [1:4, 1, 1] 0.4642 0.0206 0.9123 -1.3971
# ...
# ...


