# Author: An-Chi Ho
# Date: 13th July 2021
# Implement case 3: 6th April 2022

#---------------------------------------------------------------------
# This script shows how to use a value array as the inner dimension selector to express
# dependency on a file dimension. By this means, we don't need to specify the *_across
# parameter and Start() can recognize this dependecy relationship.
# In the first case, 'time' is dependent on 'sdate'. We create the actual time values
# for each sdate beforehand. The time array is two-dimensional with the names 'time'
# and 'sdate'.
# In the second case, 'region' is dependent on 'sdate'. The two files have different 
# index for Nino3. sdate 1993 has 'Nino3' at index 9 while sdate 2013 has 'Nino3' at 
# index 11. Create a value array for region selector so Start() can look for 'Nino3' in 
# each file.
# In the third case, 'region' is defined as an array that has dimensions 'sdate', 'member',
# and 'region'. It works if region indices is dependent on both sdate and member.
#---------------------------------------------------------------------

library(startR)
library(lubridate)

# Case 1: 'time' depends on 'sdate'
repos <- '/esarchive/exp/ecmwf/system4_m1/daily_mean/$var$_f24h/$var$_$sdate$.nc'

sdates <- ymd("20010501") + rep(years(0:2), each = 1)
times <- array(ymd("20010501") + days(0:30) + rep(years(0:2), each = 31),
               dim = c(time = 31, sdate = 3))
times <- as.POSIXct(times * 86400, tz = 'UTC', origin = '1970-01-01')

exp <- Start(dat = repos,
             var = 'tos',
             sdate = format(sdates, "%Y%m%d"),
             time = times,  #dim: [time = 31, sdate = 3]. time is corresponding to each sdate
             ensemble = indices(1:5),
             lat = 'all', 
             lon = 'all',
             synonims = list(lat = c('lat', 'latitude'),
                             lon = c('lon', 'longitude')),
             return_vars = list(lon = NULL, lat = NULL, time = 'sdate'),
             retrieve = T) 

dim(exp)
#     dat      var    sdate     time ensemble      lat      lon 
#       1        1        3       31        5      256      512 

exp[1, 1, 2, 1:10, 1, 100, 100]
# [1] 302.1276 302.1346 302.2003 302.2121 302.2552 302.3312 302.3184 302.3507
# [9] 302.3665 302.3865

summary(exp)
#    Min.  1st Qu.   Median     Mean  3rd Qu.     Max.     NA's 
#     271      274      287      287      299      305 19757385 

#=============================================================================

# Case 2: 'region' depends on 'sdate'
#NOTE: Exp "a35b" has been deleted. This example cannot be run now.
path <- paste0('/esarchive/exp/ecearth/a35b/diags/DCPP/EC-Earth-Consortium/',
               'EC-Earth3-HR/dcppA-hindcast/r1i1p1f1/Omon/$var$_mixed/gn/v20201107/',
               '$var$_Omon_EC-Earth3-HR_dcppA-hindcast_s$sdate$-r1i1p1f1_gn_$chunk$.nc')

region <- array('Nino3', dim = c(sdate = 2, region = 1))

data <- Start(dat = path,
              var = 'tosmean',
              sdate = c('1993', '2013'),
              chunk = indices(1:2),
              chunk_depends = 'sdate',
              region = region,
              time = 'all', 
              time_across = 'chunk',
              merge_across_dims = TRUE,
              return_vars = list(time = c('sdate', 'chunk'),
                                 region = 'sdate'),
              retrieve = T)

dim(data)
#   dat    var  sdate region   time 
#     1      1      2      1      2 

data[1, 1, , 1, ]
#         [,1]     [,2]
#[1,] 24.98788 24.46488  # --> region index 9 in original file
#[2,] 24.47482 24.75953  # --> region index 11 in orginal file


#=============================================================================

# Case 3: 'region' depends on 'sdate' and 'member'

#NOTE: Actually, the region indices are not dependent on sdate in this case, but
#      it should work if it is. If you have a better example, please let me know.

region <- array(c('Nino3.4', 'Nino3'), dim = c(region = 2, sdate = 3, memb = 5))

# check the array
region[, 1, 1]
#[1] "Nino3.4" "Nino3"  
region[, 1, 2]
#[1] "Nino3.4" "Nino3"  
region[, 2, 2]
#[1] "Nino3.4" "Nino3" 
#--> For each sdate-memb combination, the desired regions are "Nino3.4" and "Nino3".

path <- paste0('/esarchive/exp/ecearth/a42y/diags/DCPP/EC-Earth-Consortium/',
               'EC-Earth3/dcppA-hindcast/$memb$/Omon/$var$/gn/v*/',
               '$var$_Omon_EC-Earth3_dcppA-hindcast_s$sdate$-$memb$_gn_$chunk$.nc')

data <- Start(dat = path,
              var = 'tosmean',
              memb = paste0('r', c(24:28), 'i1p1f1'),
              region = region,
              region_var = 'region',
              sdate = paste0(2000:2002),
              time = 'all',
              chunk = 'all',
              chunk_depends = 'sdate',
              time_across = 'chunk',
              merge_across_dims = TRUE,
              return_vars = list(time = c('sdate','chunk'), 
                                 region = c('sdate', 'memb')),
              retrieve = T)

# Check output
## Nino3.4
drop(data)[ , 1, , 1]
#         [,1]     [,2]     [,3]
#[1,] 26.87246 27.28198 27.65627
#[2,] 26.87331 27.31887 27.63275
#[3,] 26.89038 27.31446 27.58801
#[4,] 26.90285 27.26750 27.66004
#[5,] 26.88851 27.28953 27.68499

## Nino3
 drop(data)[ , 2, , 1]
#         [,1]     [,2]     [,3]
#[1,] 26.58774 26.38932 26.80643
#[2,] 26.58879 26.43760 26.68655
#[3,] 26.59319 26.41373 26.64150
#[4,] 26.69607 26.40465 26.69096
#[5,] 26.59114 26.40454 26.71252