#library(startR)
#MergeArrays <- 
library(s2dverification, lib.loc = '/home/Earth/nmanuben/s2dverification/s2dverification.Rcheck')
MergeArrays <- s2dverification:::.MergeArrays

shared_dir <- 

args <- commandArgs(TRUE)[-1]

chunk_files_original <- list.files(paste0(shared_dir, '/'), '.*\\.Rds$')
chunk_files <- gsub('\\.Rds$', '', chunk_files_original)
chunk_filename_parts <- strsplit(chunk_files, '__')
if (length(unique(sapply(chunk_filename_parts, length))) != 1) {
  stop("Detected chunks with more dimensions than others.")
}
dim_names <- sapply(chunk_filename_parts[[1]], 
# TODO: strsplit by the last '_' match, not the first.
                    function(x) strsplit(x, '_')[[1]][1])
# TODO check all files have exactly the same dimnames
found_chunk_indices <- sapply(chunk_filename_parts, 
                              function(x) as.numeric(sapply(strsplit(x, '_'), '[[', 2)))
found_chunk_indices <- array(found_chunk_indices, 
                             dim = c(length(dim_names), 
                                     length(found_chunk_indices) / length(dim_names))
                            )
found_chunks_str <- apply(found_chunk_indices, 2, paste, collapse = '_')

if (length(args) > 0) {
  if ((length(args) %% 2) != 0) {
    stop("Wrong number of parameters.")
  }
  expected_dim_names <- args[(1:(length(args) / 2) - 1) * 2 + 1]
  if (any(!is.character(expected_dim_names))) {
    stop("Expected dimension names in parameters at odd positions.")
  }
  dim_indices <- args[(1:(length(args) / 2) - 1) * 2 + 2]
  if (!any(dim_indices == 'all')) {
    stop("Expected one dimension index to be 'all'.")
  }
  dim_to_merge <- which(dim_indices == 'all')
  if (length(dim_indices) > 1) {
    if (!all(is.numeric(dim_indices[-dim_to_merge]))) {
      stop("Expected all dimension index but one to be numeric.")
    }
  }
  # Check expected dim names match dim names
  ## TODO
  # Merge indices that vary along dim_to_merge whereas other fixed by dim_indices
  # REMOVE FILES
  ## TODO
  stop("Feature not implemented.")
} else {
  chunks_indices <- 1:length(dim_names)
  chunks_indices <- lapply(chunks_indices, function(x) sort(unique(found_chunk_indices[x, ])))
  names(chunks_indices) <- dim_names

  # Load all found chunks into the array 'array_of_chuks'.
  array_dims <- sapply(chunks_indices, length)
  names(array_dims) <- dim_names
  array_of_chunks <- vector('list', prod(array_dims))
  dim(array_of_chunks) <- array_dims
  array_of_chunks_indices <- array(1:prod(array_dims), array_dims)
  for (i in 1:prod(array_dims)) {
    chunk_indices <- which(array_of_chunks_indices == i, arr.ind = TRUE)[1, ]
    j <- 1
    chunk_indices_on_file <- sapply(chunk_indices,
                                    function(x) {
                                      r <- chunks_indices[[j]][x]
                                      j <<- j + 1
                                      r
                                     })
    found_chunk <- which(found_chunks_str == paste(chunk_indices_on_file, 
                                                   collapse = '_'))[1]
    if (length(found_chunk) > 0) {
      array_of_chunks[[i]] <- readRDS(chunk_files_original[found_chunk])
    }
  }

  # Merge the chunks.
  for (dim_index in 1:length(dim_names)) {
    dim_sub_array_of_chunks <- dim_sub_array_of_chunk_indices <- NULL
    if (dim_index < length(dim_names)) {
      dim_sub_array_of_chunks <- sapply(chunks_indices[(dim_index + 1):length(dim_names)],
                                        length)
      names(dim_sub_array_of_chunks) <- dim_names[(dim_index + 1):length(dim_names)]
      dim_sub_array_of_chunk_indices <- dim_sub_array_of_chunks
      sub_array_of_chunk_indices <- array(1:prod(dim_sub_array_of_chunk_indices), 
                                          dim_sub_array_of_chunk_indices)
    } else {
      sub_array_of_chunk_indices <- NULL
    }
    sub_array_of_chunks <- vector('list', prod(dim_sub_array_of_chunks))
    dim(sub_array_of_chunks) <- dim_sub_array_of_chunks
    for (i in 1:prod(dim_sub_array_of_chunks)) {
      if (!is.null(sub_array_of_chunk_indices)) {
        chunk_sub_indices <- which(sub_array_of_chunk_indices == i, arr.ind = TRUE)[1, ]
      } else {
        chunk_sub_indices <- NULL
      }
      for (j in 1:(array_dims[dim_index])) {
        new_chunk <- do.call('[[', c(list(x = array_of_chunks), 
                                     as.list(c(j, chunk_sub_indices))))
        #do.call('[[<-', c(list(x = array_of_chunks), 
        #                  as.list(c(j, chunk_sub_indices)),
        #                  list(value = NULL)))
        if (is.null(new_chunk)) {
          stop("Chunks missing.")
        }
        if (is.null(sub_array_of_chunks[[i]])) {
          sub_array_of_chunks[[i]] <- list(mod = NULL, obs = NULL)
        }
        if (is.null(new_chunk$mod)) {
          if (!is.null(sub_array_of_chunks[[i]]$mod)) {
            stop("Chunks missing (model data).")
          }
        } else if (is.null(sub_array_of_chunks[[i]]$mod)) {
          sub_array_of_chunks[[i]]$mod <- new_chunk$mod
        } else {
          sub_array_of_chunks[[i]]$mod <- MergeArrays(sub_array_of_chunks[[i]]$mod, 
                                                      new_chunk$mod, 
                                                      dim_names[dim_index])
        }
        if (is.null(new_chunk$obs)) {
          if (!is.null(sub_array_of_chunks[[i]]$obs)) {
            stop("Chunks missing (obs data).")
          }
        } else if (is.null(sub_array_of_chunks[[i]]$obs)) {
          sub_array_of_chunks[[i]]$obs <- new_chunk$obs
        } else {
          sub_array_of_chunks[[i]]$obs <- MergeArrays(sub_array_of_chunks[[i]]$obs, 
                                                      new_chunk$obs, 
                                                      dim_names[dim_index])
        }
      }
    }
    array_of_chunks <- sub_array_of_chunks
    rm(sub_array_of_chunks)
    gc()
  }
  result <- array_of_chunks[[1]] 
}

sapply(chunk_files_original, file.remove)

saveRDS(result, file = paste0(shared_dir, '/result.Rds'))
