context("CSTools::CST_RFSlope tests")

##############################################
# dat1
data <- rnorm(2 * 10 * 4 * 5 * 6 * 7)
dim(data) <- c(dataset = 2, member = 10,
                sdate = 4, ftime = 5, lat = 6, lon = 7)
lon <- seq(0, 12, 2)
lat <- seq(10, 15, 1)
coords <- list(lon = lon, lat = lat)
exp <- list(data = data, coords = coords)
attr(exp, "class") <- "s2dv_cube"

# dat2
exp2 <- exp
exp2$attrs <- list(source_files = 'exp')
exp2$coords <- NULL
exp2_2 <- exp2
exp2_2$coords <- list(long = seq(1:4), lati = seq(1:4))
exp2_3 <- exp
names(dim(exp2_3$data)) <- c("dataset", "member", "sdate", "ftime", "lati", "loni")

##############################################
test_that("1. Input checks", {
  # Check 's2dv_cube'
  expect_error(
    CST_RainFARM(rnorm(2 * 15 * 4 * 5 * 6 * 7)),
    paste0("Parameter 'data' must be of the class 's2dv_cube', ",
           "as output by CSTools::CST_Load.")
  )
  # Check 'exp' object structure
  expect_error(
    CST_RainFARM(exp2),
    paste0("Parameter 'data' must have 'data' and 'coords' elements ",
           "within the 's2dv_cube' structure.")
  )
  # Check coordinates
  expect_error(
    CST_RainFARM(exp2_2),
    paste0("Spatial coordinate names do not match any of the names accepted by ",
           "the package.")
  )
  # Check dimensions
  expect_error(
    CST_RainFARM(exp2_3),
    paste0("Spatial dimension names do not match any of the names accepted by ",
           "the package.")
  )
  # Generate simple synthetic data
  # 4x5 in space, 2 members, 3 sdates, 6 ftime
  r <- exp(rnorm(1 * 2 * 3 * 6 * 5 * 4))
  dim(r) <- c(dataset = 1, member = 2, sdate = 3, ftime = 6, lat = 5, lon = 4)
  lon <- seq(0, 6, 2)
  lat <- seq(10, 18, 2)
  exp <- list(data = r, coords = list(lat = lat, lon = lon))
  attr(exp, 'class') <- 's2dv_cube'

  expect_warning(
    res <- CST_RainFARM(exp, nf=8),
    "input data are expected to be on a square grid with an even number of pixels per side."
  )
  # Using a correct square dataset
  r <- exp(rnorm(1 * 2 * 3 * 6 * 4 * 4))
  dim(r) <- c(dataset = 1, member = 2, sdate = 3, ftime = 6, lat = 4, lon = 4)
  lat <- seq(10, 16, 2)
  exp <- list(data = r, coords = list(lat = lat, lon = lon))
  attr(exp, 'class') <- 's2dv_cube'

  expect_warning(
    res <- CST_RainFARM(exp, nf=8),
    "Selected time dim: ftime"
  )
  expect_error(
    res <- CST_RainFARM(exp, nf=8, weights=array(0,dim=c(2,2))),
    paste0("Parameters 'lon_dim' and 'lat_dim' do not match with 'weights' ", 
           "dimension names.")
  )
})

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

test_that("2. Simple use case", {
  r <- exp(rnorm(1 * 2 * 3 * 6 * 4 * 4))
  dim(r) <- c(dataset = 1, member = 2, sdate = 3, ftime = 6, lat = 4, lon = 4)
  lat <- seq(10, 16, 2)
  lon <- seq(0, 6, 2)
  exp <- list(data = r, coords = list(lat = lat, lon = lon))
  attr(exp, 'class') <- 's2dv_cube'
  
  dimexp = dim(exp$data)

  res <- CST_RainFARM(exp, nf = 8, time_dim = c("ftime", "sdate"), 
                      slope = 1.7, nens = 2)
  expect_equal(dim(res$data), c(dimexp["dataset"], dimexp["member"], 
                                realization = 2, dimexp["sdate"],
                                dimexp["ftime"], dimexp["lat"] * 8,
                                dimexp["lon"] * 8))

  expect_equivalent(length(res$coords$lon), dimexp["lon"] * 8) 
  expect_equivalent(length(res$coords$lat), dimexp["lat"] * 8) 

  res <- CST_RainFARM(exp, nf = 8, time_dim = c("ftime", "sdate"),
                      nens = 2, drop_realization_dim = TRUE)
  expect_equal(dim(res$data), c(dimexp["dataset"], dimexp["member"] * 2,
                                dimexp["sdate"],
                                dimexp["ftime"], dimexp["lat"] * 8,
                                dimexp["lon"] * 8))

  res <- CST_RainFARM(exp, nf = 8, time_dim = c("ftime", "sdate"), slope = 1.7,
                      nens = 2, nproc = 2, fsmooth = FALSE)
  expect_equal(dim(res$data), c(dimexp["dataset"], dimexp["member"], 
                                realization = 2, dimexp["sdate"],
                                dimexp["ftime"], dimexp["lat"] * 8,
                                dimexp["lon"] * 8))

  # Use rainfarmr package functions to test for consistency

  # Aggregated downscaled data should be the same as original

  expect_equivalent(agg(res$data[1,1,1,1,1,,], 4), exp$data[1,1,1,1,,])

  res <- CST_RainFARM(exp, nf = 8, time_dim = c("ftime", "sdate"),
                      nens = 2, nproc = 2, fglob = TRUE)

  expect_equal(mean(agg(res$data[1,1,1,1,1,,], 4)),
               mean(exp$data[1,1,1,1,,]))

  # Create a more realistic perfect-model precipitation
  z <- 1 : (32 * 32)
  dim(z) <- c(32, 32)
  x=rep(-15:16,32); dim(x) <- c(32, 32);  x=x+4 ; y<-t(x);
  z<-sqrt(x^2+y^2); z=(z<=12)*(z>=8)*0.8+0.6
  
  exppf <- exp 
  nens = 10
  exppf$data <- array(1. , c(1, nens, 3, 6, 4, 4))
  dim(exppf$data) <- c(dataset = 1, member = nens, sdate = 3,
                       ftime = 6, lat = 4, lon = 4)
  rpf <- array(1. , c(1, nens, 3, 6, 32, 32))
  for (i in 1:nens) {
    for (j in 1:3) {
      for (k in 1:6) {
        f = initmetagauss(1.7, 32)
        r = exp(metagauss(f)) * z
        rpf[1, i, j, k, , ] <- r
        exppf$data[1, i, j, k, , ] <- agg(r, 4)
      }
    }
  }
  rpfm=agg(apply(rpf, c(5, 6), mean),32)

  # Use climatological mean of PF precipitation to generate sythetic weights
  w <- rfweights(rpfm, res$coords$lon, res$coords$lat, exp$coords$lon, 
                 exp$coords$lat, 8, fsmooth = FALSE )
  names(dim(w)) <- c('lon', 'lat')
  res <- CST_RainFARM(exppf, nf = 8, time_dim = c("ftime", "sdate", "member"),
                      nens = 2, nproc = 2, fsmooth = FALSE)
  resw <- CST_RainFARM(exppf, nf = 8, time_dim = c("ftime", "sdate", "member"),
                       nens = 2, nproc = 2, fsmooth = FALSE, weights = w)
  resm <- agg(apply(res$data, c(6,7), mean),32)
  reswm <- agg(apply(resw$data, c(6,7), mean),32)
  
  # Things should work much  better with weights compared to no weights
  expect_lt(mean((reswm-rpfm)^2), mean((resm-rpfm)^2)*0.5)

  # Test if downscaled data have similar slope to input data
  # Synthetic fine-scale data with slope 1.7 
  f = initmetagauss(1.7, 256)
  r = exp(metagauss(f))
  ra <- agg(r, 8)
  dim(ra) <- c(dataset = 1, member = 1, sdate = 1, ftime = 1, lat = 8, lon = 8)
  expcoarse <- exp
  expcoarse$data <- ra
  dim(r) <- c(dataset = 1, member = 1, sdate = 1, ftime = 1, lat = 256, lon = 256)
  expfine <- exp
  expfine$data <- r
  res <- CST_RainFARM(expcoarse, nf = 32, time_dim=c("ftime", "sdate"),
                      slope = 1.7, fsmooth=FALSE, drop_realization_dim=TRUE)
  # TO DO: Develop within the new s2dv_cube
  
  sres= CST_RFSlope(res, time_dim = c("ftime", "sdate"))
  sexp= CST_RFSlope(expfine, time_dim = c("ftime", "sdate"))
  expect_equal(sres, sexp, tolerance=0.25)
})  
