#Define a function that receives 3 multidimensional arrays:
#     - dat1: One with the dimensions 'time', 'lat' and 'lon'.
#     - dat2: One with the dimensions 'lat' and 'lon'.
#     - dat3: One with the dimensions 'time' and 'lat'.
#Plus the following additional parameters:
#     - a Boolean parameter called revert_order
#     - the '...' parameter

#The function should:
#    a) take the second array and subtract it to each time step in the first array.
#    b) for each time step in the first array, all values along the 'lon' dimension for each 'lat' should be multiplied by the corresponding value in the third array.
#    c) for each time step, the area average is computed (using the mean() function, to which the additional parameters in '...' are forwarded).
#    d) if revert_order == TRUE, b) should be performed before a).

#The function should return an array with the dimension 'time'.

#-------
#dat1 <- 1 : 100
#dim(dat1) <- c(time = 10, lat = 5, lon = 2)
#
#dat2 <- rep(1:5, 2)
#dim(dat2) <- c(lat = 5, lon = 2)
#
#dat3 <- rep(1:5, each = 10)
#dim(dat3) <- c(time = 10, lat = 5)
#
#weighted <- c(1, 0.8, 0.6, 0.4, 0.2)
#
#--------
library(plyr)

func_apply <- function(dat1, dat2, dat3, revert_order = FALSE, weighted = NULL, ...) {
  # Check dat1
  if (!is.numeric(dat1) || !is.array(dat1)) {
    stop("Parameter 'dat1' must be a numeric array.")
  }
  input_error <- FALSE
  if (length(dim(dat1)) != 3) {
    input_error <- TRUE
  }
  if (is.null(names(dim(dat1)))) {
    input_error <- TRUE
  }
  if (!identical(names(dim(dat1)), c('time', 'lat', 'lon'))) {
    input_error <- TRUE
  }
  if (input_error) {
    stop("Parameter 'dat1' must have the following named dimensions, in the ",
         "following order: 'time', 'lat', 'lon'.")
  }

  # Check dat2
  if (!is.numeric(dat2) || !is.array(dat2)) {
    stop("Parameter 'dat2' must be a numeric array.")
  }
  input_error <- FALSE
  if (length(dim(dat2)) != 2) {
    input_error <- TRUE
  }
  if (is.null(names(dim(dat2)))) {
    input_error <- TRUE
  }
  if (!identical(names(dim(dat2)), c('lat', 'lon'))) {
    input_error <- TRUE
  }
  if (input_error) {
    stop("Parameter 'dat2' must have the following named dimensions, in the ",
         "following order: 'lat', 'lon'.")
  }

  # Check dat3
  if (!is.numeric(dat3) || !is.array(dat3)) {
    stop("Parameter 'dat3' must be a numeric array.")
  }
  input_error <- FALSE
  if (length(dim(dat3)) != 2) {
    input_error <- TRUE
  }
  if (is.null(names(dim(dat3)))) {
    input_error <- TRUE
  }
  if (!identical(names(dim(dat3)), c('time', 'lat'))) {
    input_error <- TRUE
  }
  if (input_error) {
    stop("Parameter 'dat3' must have the following named dimensions, in the ",
         "following order: 'time', 'lat'.")
  }

  # Consistency checks dat1, dat2, dat3
  time_dim_len <- dim(dat1)['time']
  lat_dim_len <- dim(dat1)['lat']
  lon_dim_len <- dim(dat1)['lon']
  if (dim(dat1)['time'] != time_dim_len || dim(dat3)['time'] != time_dim_len) {
    stop("Found an inconsistency in the length of the dimension 'time' in ",
         "parameters 'dat1' and 'dat3'.")
  }
  if (dim(dat1)['lat'] != lat_dim_len || dim(dat2)['lat'] != lat_dim_len ||
      dim(dat3)['lat'] != lat_dim_len) {
    stop("Found an inconsistency in the length of the dimension 'lat' in ",
         "parameters 'dat1', 'dat2' and 'dat3'.")
  }
  if ((dim(dat1)['lon'] != lon_dim_len) || (dim(dat2)['lon'] != lon_dim_len)) {
    stop("Found an inconsistency in the length of the dimension 'lon' in ",
         "parameters 'dat1', 'dat2' and 'dat3'.")
  }

  # Check revert_order
  if (!is.logical(revert_order)) {
    stop("Parameter 'revert_order' must be a logical value.")
  }
  if (length(revert_order) != 1) {
    stop("Parameter 'revert_order' must be of length 1.")
  }

  dim1 <- dim(dat1)
  dim2 <- dim(dat2)
  dim3 <- dim(dat3)

  if (revert_order == FALSE) {
    dat2_rep <- array(dat2, dim = c(dim2, dim1['time']))
    dat2_rep <- s2dverification:::.aperm2(dat2_rep, c(3, 1, 2))

    res1 <- dat1 - dat2_rep

    dat3_rep <- array(dat3, dim = c(dim3, dim(res1)['lon']))
    res2 <- res1 * dat3_rep
  } else {
    dat3_rep <- array(dat3, dim = c(dim3, dim1['lon']))
 
    res1 <- dat1 * dat3_rep

    dat2_rep <- array(dat2, dim = c(dim2, dim(res1)['time']))
    dat2_rep <- s2dverification:::.aperm2(dat2_rep, c(3, 1, 2))

    res2 <- res1 - dat2_rep
  }

#c)
  #weighted? need lat/weighted as an input

  fun_wt <- function(x, ...) {
    weighted.mean(x, weighted, ...)
  }

  f1 <- function(res2, weighted, ...) {
    if (is.null(weighted)) {
      result <- apply(res2, which(names(dim(res2)) == 'time'), mean, ...)
    } else {
      tmp <- aaply(res2, c(which(names(dim(res2)) == 'time'), which(names(dim(res2)) == 'lat')), 
                   mean, ...)
      result <- apply(tmp, 1, FUN = fun_wt, ...)
    }
    return(result)
  }

  res3 <- f1(res2, weighted = weighted, ...)  
  dim(res3) <- c(time = length(res3))

  return(res3)
}
