context("CSTools::CST_Subset tests")

##############################################

library(startR)

##############################################

test_that("1. Input checks: CST_Subset", {
  # Check that x is s2dv_cube
  expect_error(
    CST_Subset(array(10)), 
    "Parameter 'x' must be of the class 's2dv_cube'."
  )
  # Check var_dim
  expect_error(
    CST_Subset(lonlat_prec, var_dim = 1),
    "Parameter 'var_dim' must be a character string."
  )
  expect_error(
    CST_Subset(lonlat_prec, var_dim = c('tas', 'psl')),
    "Parameter 'var_dim' must be a character string."
  )
  # Check dat_dim
  expect_error(
    CST_Subset(lonlat_prec, dat_dim = 1),
    "Parameter 'dat_dim' must be a character string."
  )
  expect_error(
    CST_Subset(lonlat_prec, dat_dim = c('dat1', 'dat2')),
    "Parameter 'dat_dim' must be a character string."
  )
})

##############################################

test_that("2. Output checks: CST_Subset", {
  res1 <- CST_Subset(lonlat_prec, along = c('lat', 'lon', 'sdate', 'member'),
                     indices = list(1, 1:2, 1, 1:2),
                     drop = 'all')
  # Check dimensions
  expect_equal(
    dim(res1$data), 
    res1$dims
  )
  expect_equal(
    dim(res1$data), 
    c(member = 2, ftime  = 31, lon = 2)
  )
  # Check coordinates
  expect_equal(
    names(res1$coords),
    c("member", "ftime", "lon")
  )
  # Check attrs
  expect_equal(
    names(res1$attrs),
    names(lonlat_prec$attrs)
  )
  expect_equal(
    names(res1$attrs$Variable$metadata),
    c("lon", "prlr")
  )
  expect_equal(
    res1$attrs$Datasets,
    c("exp1")
  )
  # Check 'dat_dim'
  res2 <- CST_Subset(lonlat_prec, along = c('lat'), indices = list(1), 
                     drop = 'all', dat_dim = 'dataset')
  res3 <- CST_Subset(lonlat_prec, along = c('lat'), indices = list(1), 
                     drop = 'selected', dat_dim = 'dataset')
  res4 <- CST_Subset(lonlat_prec, along = c('dataset'), indices = list(1), 
                     drop = 'all', dat_dim = 'dataset')
  res5 <- CST_Subset(lonlat_prec, along = c('dataset'), indices = list(1), 
                     drop = 'selected', dat_dim = 'dataset')
  expect_equal(
    res2$attrs$Datasets,
    res3$attrs$Datasets
  )
  expect_equal(
    length(res4$attrs$Datasets),
    length(res5$attrs$Datasets)
  )
  # Check 'Dates'
  res6 <- CST_Subset(lonlat_prec, along = c('sdate', 'ftime'), 
                     indices = list(1, 1:10), drop = 'selected')
  res7 <- CST_Subset(lonlat_prec, along = c('sdate', 'ftime'), 
                     indices = list(1, 1:10), drop = 'none')
  # Dates dimensions
  expect_equal(
    dim(res6$attrs$Dates),
    res6$dims[which(names(dim(res6$data)) %in% c('sdate', 'ftime'))]
  )
  expect_equal(
    dim(res7$attrs$Dates),
    c(ftime = 10, sdate = 1)
  )
  # sdates coordinates
  expect_equal(
    names(res6$coords),
    c("dataset", "member", "ftime", "lat", "lon")
  )
  expect_equal(
    as.vector(res7$coords$sdate),
    c("20101101")
  )
})

##############################################

repos1 <- "/esarchive/exp/ecmwf/system5_m1/monthly_mean/$var$_f6h/$var$_$sdate$.nc"
repos2 <- "/esarchive/exp/ecmwf/system4_m1/monthly_mean/$var$_f6h/$var$_$sdate$.nc"

suppressWarnings(
  data <- Start(dat = list(list(name = 'system4_m1', path = repos2),
                            list(name = 'system5_m1', path = repos1)),
                var = c('tas', 'sfcWind'),
                sdate = c('20170101'),
                ensemble = indices(1),
                time = indices(1),
                lat = indices(1:10),
                lon = indices(1:10),
                synonims = list(lat = c('lat', 'latitude'),
                                lon = c('lon', 'longitude')),
                return_vars =  list(time = 'sdate',
                                    longitude = 'dat',
                                    latitude = 'dat'),
                metadata_dims = c('dat', 'var'),
                retrieve = T)
)
suppressWarnings(
  exp_start <- as.s2dv_cube(data)
)

##############################################

test_that("3. Output checks with Start", {
  res8 <- CST_Subset(exp_start, along = c('lon', 'sdate', 'ensemble', 'var', 'dat'),
                     indices = list(1:2, 1, 1, 1, 1),
                     drop = 'none')
  res9 <- CST_Subset(exp_start, along = c('lon', 'sdate', 'ensemble', 'var', 'dat'),
                     indices = list(1:2, 1, 1, 1, 1),
                     drop = FALSE, var_dim = 'var', dat_dim = 'dat')
  res10 <- CST_Subset(exp_start, along = c('lon', 'sdate', 'ensemble', 'var', 'dat'),
                     indices = list(1:2, 1, 1, 1, 1),
                     drop = 'selected', var_dim = 'var', dat_dim = 'dat')
  # Check dimensions
  expect_equal(
    dim(res8$data), 
    c(dat = 1, var = 1, sdate = 1, ensemble = 1, time = 1, lat = 10, lon = 2)
  )
  expect_equal(
    dim(res8$data), 
    dim(res9$data)
  )
  expect_equal(
    dim(res10$data), 
    c(time = 1, lat = 10, lon = 2)
  )
  # Check coordinates
  expect_equal(
    names(res8$coords),
    names(res8$dims)
  )
  expect_equal(
    names(res9$coords),
    names(res9$dims)
  )
  # varName
  expect_equal(
    res8$attrs$Variable$varName,
    c("tas", "sfcWind")
  )
  expect_equal(
    res9$attrs$Variable$varName,
    c("tas")
  )
  expect_equal(
    res10$attrs$Variable$varName,
    NULL
  )
  # metadata
  expect_equal(
    names(res8$attrs$Variable$metadata),
    c("time", "lat", "lon", "tas", "sfcWind")
  )
  expect_equal(
    names(res9$attrs$Variable$metadata),
    c("time", "lat", "lon", "tas")
  )
  expect_equal(
    names(res10$attrs$Variable$metadata),
    c("time", "lat", "lon")
  )
  # Datasets
  expect_equal(
    res8$attrs$Datasets,
    c("system4_m1", "system5_m1")
  )
  expect_equal(
    res9$attrs$Datasets,
    c("system4_m1")
  )
  expect_equal(
    length(res10$attrs$Datasets),
    0
  )
  # Check source_files
  expect_equal(
    dim(res8$attrs$source_files),
    c(dat = 1, var = 1, sdate = 1)
  )
  expect_equal(
    dim(res9$attrs$source_files),
    c(dat = 1, var = 1, sdate = 1)
  )
  expect_equal(
    dim(res10$attrs$source_files),
    c(1) 
  )
})