# This unit test is to check chunking over the split dim. It involves 
# how to arrange the chunks in a correct order even when chunking is happening.

context("Chunk over split dim")

test_that("1. The files are not repeated", {

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')

suppressWarnings(
exp <- Start(dat = repos_exp,
             var = 'tas',
             sdate = as.character(c(2005:2008)),
             time = indices(1:3),
             lat = indices(1:2),
             lon = indices(1:3),
             synonims = list(lat = c('lat', 'latitude'),
                             lon = c('lon', 'longitude')),
             return_vars = list(lon = NULL,
                                lat = NULL,
                                time = 'sdate'),
             retrieve = FALSE)
)
lats <- attr(exp, 'Variables')$common$lat
lons <- attr(exp, 'Variables')$common$lon
## The 'time' attribute is a two-dim array
dates <- attr(exp, 'Variables')$common$time
#dim(dates)
#sdate  time 
#    4     3

repos_obs <- '/esarchive/recon/ecmwf/erainterim/monthly_mean/$var$_f6h/$var$_$date$.nc'
suppressWarnings(
obs <- Start(dat = repos_obs,
             var = 'tas',
             date = unique(format(dates, '%Y%m')),
             time = values(dates),  #dim: [sdate = 4, time = 3]
             lat = values(lats),
             lat_reorder = Sort(),
             lon = values(lons),
             lon_reorder = CircularSort(0, 360),
             time_across = 'date',
             merge_across_dims = TRUE,
             split_multiselected_dims = TRUE,
             synonims = list(lat = c('lat', 'latitude'),
                             lon = c('lon', 'longitude')),
             return_vars = list(lon = NULL,
                                lat = NULL,
                                time = 'date'),
             retrieve = FALSE)
)

fun <- function(x) {
return(x)
}
step1 <- Step(fun, target_dims = 'time', output_dims = 'time')
wf1 <- AddStep(obs, step1)
suppressWarnings(
res <- Compute(wf1, chunks = list(sdate = 1))$output1
)
suppressWarnings(
res1 <- Compute(wf1, chunks = list(sdate = 2))$output1
)

step2 <- Step(fun, target_dims = 'sdate', output_dims = 'sdate')
wf2 <- AddStep(obs, step2)
suppressWarnings(
res2 <- Compute(wf2, chunks = list(time = 2))$output1
)

step3 <- Step(fun, target_dims = 'lon', output_dims = 'lon')
wf3 <- AddStep(obs, step3)
suppressWarnings(
res3 <- Compute(wf3, chunks = list(time = 2, sdate = 2))$output1
)

expect_equal(
res[1,1,1,,2,2],
c(250.8127, 248.7926, 249.4770, 246.8312),
tolerance = 0.0001
)
expect_equal(
res[2,1,1,,2,2],
c(237.3297, 235.9947, 235.4366, 235.9104),
tolerance = 0.0001
)
expect_equal(
res,
res1
)
expect_equal(
res,
aperm(res2, c(4,2,3,1,5,6))
)
expect_equal(
res,
aperm(res3, c(5,2,3,4,6,1))
)
expect_equal(
dim(res1),
c(time = 3, dat = 1, var = 1, sdate = 4, lat = 2, lon = 3)
)
expect_equal(
dim(res2),
c(sdate = 4, dat = 1, var = 1,  time = 3, lat = 2, lon = 3)
)
expect_equal(
dim(res3),
c(lon = 3, dat = 1, var = 1, sdate = 4, time = 3, lat = 2)
)

})


test_that("2. The files are repeated", {

ecmwf_path_hc <- paste0('/esarchive/exp/ecmwf/s2s-monthly_ensforhc/daily/$var$_f24h/$sdate$/$var$_$syear$.nc')
obs_path <-paste0('/esarchive/recon/ecmwf/era5/daily/$var$-r240x121/$var$_$file_date$.nc')
sdates.seq <- c("20161222","20161229", "20170105","20170112")

suppressWarnings(
hcst <- Start(dat = ecmwf_path_hc,
              var = 'sfcWind',
              sdate = sdates.seq,
              syear =  indices(1:2), #'all',
              time = 'all',
              latitude = indices(1),
              longitude = indices(1),
              ensemble = indices(1),
              syear_depends = 'sdate',
              return_vars = list(latitude = 'dat',
                                 longitude = 'dat',
                                 time = c('sdate','syear')
                                  ),
              retrieve = F)
)
datess <- attr(hcst, 'Variables')$common$time
file_date <- unique(sapply(datess, format, '%Y%m'))

suppressWarnings(
obs <- Start(dat = obs_path,
             var = 'windagl100',
             latitude = indices(1),
             longitude = indices(1:2),
             file_date= file_date,
             time = values(datess),    # 'sdate' 'syear' 'time'
             time_var = 'time',
             time_across = 'file_date',
             merge_across_dims= TRUE,
             merge_across_dims_narm = TRUE,
             split_multiselected_dims = TRUE,
             synonims = list(latitude=c('lat','latitude'),
                             longitude=c('lon','longitude')),
             return_vars = list(latitude = 'dat',
                                longitude = 'dat',#
                                time = c('file_date')),
             retrieve = F)
)

fun <- function(x) {
return(x)
}
st1 <- Step(fun, target_dims = 'time', output_dims = 'time')
adst1 <- AddStep(obs, st1)
suppressWarnings(
res <- Compute(adst1, chunks = list(sdate = 1))$output1
)
suppressWarnings(
res1 <- Compute(adst1, chunks = list(sdate = 2))$output1
)

st2 <- Step(fun, target_dims = 'sdate', output_dims = 'sdate')
adst2 <- AddStep(obs, st2)
suppressWarnings(
res2 <- Compute(adst2, chunks = list(time = 2))$output1
)

st3 <- Step(fun, target_dims = 'longitude', output_dims = 'longitude')
adst3 <- AddStep(obs, st3)
suppressWarnings(
res3 <- Compute(adst3, chunks = list(time = 2, sdate = 2))$output1
)

expect_equal(
drop(res)[1,1,,2],
c(8.324215, 10.622266, 14.304168, 12.299168),
tolerance = 0.0001
)
expect_equal(
drop(res)[2,1,,1],
c(4.357206, 21.909714, 4.808544, 6.861799),
tolerance = 0.0001
)
expect_equal(
res,
res1
)
expect_equal(
res,
aperm(res2, c(7, 2, 3, 4, 5, 1, 6))
)
expect_equal(
res,
aperm(res3, c(7, 2, 3, 4, 1, 5, 6))
)
expect_equal(
dim(res1),
c(time = 47, dat = 1, var = 1, latitude = 1, longitude = 2, sdate = 4, syear = 2)
)
expect_equal(
dim(res2),
c(sdate = 4, dat = 1, var = 1, latitude = 1, longitude = 2, syear = 2, time = 47)
)
expect_equal(
dim(res3),
c(longitude = 2, dat = 1, var = 1, latitude = 1, sdate = 4, syear = 2, time = 47)
)


})
