context("Compute, extra function arguments")

test_that("ex2_6", {


#=========================
# Prepare sdates and paths
#=========================
  dataset <- "/esarchive/recon/ecmwf/era5/daily_mean/$var$_f1h/$var$_$sdate$.nc"
  sdates <- paste0(1981:1982, rep(10:12, 2))
#===================
# Get daily winds 
#===================
suppressWarnings(
  wind <- Start(dataset = dataset,
                  var = "sfcWind",
                  sdate = sdates,
                  time = 'all',
                  longitude = indices(1:3),
                  latitude = indices(1:2),
                  return_vars = list(time = NULL, latitude = NULL, longitude = NULL),
                  retrieve = FALSE,
                  synonims = list(longitude = c('lon', 'longitude'),
                                  latitude = c('lat', 'latitude')))
)
  
  # Synthetic MJO for 'season = OND':
  set.seed(1)
  MJO <- data.frame(vdate = 1:(30 * 2 + 31 * 4),
                    phase = c(rep(1:8, 23)),
                    amplitude = 10 * rnorm(31 * 4 + 30 * 2))
  
  stratify_atomic <- function(field, MJO, season = c("JFM", "OND"),
                              lag = 0, ampl = 2, relative = TRUE, signif = 0.05) {
      # Arrange wind in form (days) to match MJO
      nmonths <- dim(field)[3]
      field <- aperm(field, c(1, 2, 4, 3))
      dim(field) <- c(31 * nmonths)
      if(season == "JFM") {
          daysok <- rep(c(rep(TRUE, 31), rep(TRUE, 28),
                          rep(FALSE, 3), rep(TRUE, 31)), nmonths / 3)
      } else if (season == "OND") {
          daysok <- rep(c(rep(TRUE, 31), rep(TRUE, 30),
                          rep(FALSE, 1), rep(TRUE, 31)), nmonths / 3)
      }
      field <- field[daysok]
      dim(field) <- c(days = length(field))
  
      if(dim(field)[1] != dim(MJO)[1]) {
          stop("MJO indices and wind data have different number of days")
      }
  
      idx <- function(MJO, phase, ampl, lag){
          if(lag == 0) {
              return(MJO$phase == phase & MJO$amplitude > ampl)
          }
          if(lag > 0) {
              return(dplyr::lag(MJO$phase == phase & MJO$amplitude > ampl,
                                lag, default = FALSE))
          }
          if(lag < 0) {
              return(dplyr::lead(MJO$phase == phase & MJO$amplitude > ampl,
                                - 1 * lag, default = FALSE))
          }
      }
      strat <- plyr::laply(1:8, function(i) {
                               idx2 <- idx(MJO, i, ampl, lag)
                               if (relative) {
                                     return(mean(field[idx2]) / mean(field) - 1)
                               } else {
                                     return(mean(field[idx2]) - mean(field))
                               }})
      strat.t.test <- plyr::laply(1:8, function(i) {
                              idx2 <- idx(MJO, i, ampl, lag)
                              return(t.test(field[idx2], field)$p.value)})
      return(list(strat = strat, t_test = strat.t.test))
  }
  
  step <- Step(stratify_atomic,
               target_dims = list(field = c('dataset', 'var', 'sdate', 'time')),
               output_dims = list(strat = c('phase'), t_test = c('phase')))
  workflow <- AddStep(wind, step, MJO = MJO, season = "OND", lag = "0", amp = 0)

suppressWarnings(  
  res <- Compute(workflow$strat,
                 chunks = list(latitude = 2))
)

expect_equal(
attr(wind, 'Dimensions'),
c(dataset = 1, var = 1, sdate = 6, time = 31, longitude = 3, latitude = 2)
)
expect_equal(
names(res),
c('strat', 't_test')
)
expect_equal(
dim(res$strat),
c(phase = 8, longitude = 3, latitude = 2)
)
expect_equal(
dim(res$t_test),
c(phase = 8, longitude = 3, latitude = 2)
)
expect_equal(
mean(res$strat),
-0.01373227,
tolerance = 0.0001
)
expect_equal(
res$strat[1:6,2,1],
c(-0.002499522, 0.125437301, -0.044554040, -0.034862961, 0.019349007, -0.143963809),
tolerance = 0.0001
)
expect_equal(
res$t_test[1:6,2,1],
c(0.9808923, 0.3378701, 0.6251017, 0.7305827, 0.8573760, 0.2473257),
tolerance = 0.0001
)
expect_equal(
mean(res$t_test),
0.6419336,
tolerance = 0.0001
)

})
