#' @rdname CST_RainFARM
#' @title RainFARM stochastic precipitation downscaling of a CSTools object
#'
#' @author Jost von Hardenberg - ISAC-CNR, \email{j.vonhardenberg@isac.cnr.it}
#'
#' @description This function implements the RainFARM stochastic precipitation
#' downscaling method and accepts a CSTools object (an object of the class 
#' 's2dv_cube' as provided by `CST_Load`) as input.
#' Adapted for climate downscaling and including orographic correction
#' as described in Terzago et al. 2018.
#' References: Terzago, S. et al. (2018). NHESS 18(11), 2825–2840.
#' http://doi.org/10.5194/nhess-18-2825-2018 ,
#' D'Onofrio et al. (2014), J of Hydrometeorology 15, 830-843; Rebora et. al. (2006), JHM 7, 724.
#' Please notice: The function uses the \code{julia_setup()} function from the JuliaCall library, 
#' which takes a long time only the first time it is called on a machine.
#'
#' @param data An object of the class 's2dv_cube' as returned by `CST_Load`, 
#' containing the spatial precipitation fields to downscale.
#' The data object is expected to have an element named \code{$data} with at least two
#' spatial dimensions named "lon" and "lat" and one or more dimensions over which
#' to compute average spectral slopes (unless specified with parameter \code{slope}),
#' which can be specified by parameter \code{time_dim}.
#' The number of longitudes and latitudes in the input data is expected to be even and the same. If not
#' the function will perform a subsetting to ensure this condition.
#' @param weights Matrix with climatological weights which can be obtained using
#' the \code{RFWeights} function. If \code{weights=1.} (default) no weights are used.
#' The matrix should have dimensions (lon, lat) in this order.
#' The names of these dimensions are not checked. 
#' @param nf Refinement factor for downscaling (the output resolution is increased by this factor).
#' @param slope Prescribed spectral slope. The default is \code{slope=0.}
#' meaning that the slope is determined automatically over the dimensions specified by \code{time_dim}.
#' @param kmin First wavenumber for spectral slope (default: \code{kmin=1}).
#' @param nens Number of ensemble members to produce (default: \code{nens=1}).
#' @param fglob Logical to conserve global precipitation over the domain (default: FALSE).
#' @param fsmooth Logical to conserve precipitation with a smoothing kernel (default: TRUE).
#' @param time_dim String or character array with name(s) of dimension(s)
#' (e.g. "ftime", "sdate", "member" ...) over which to compute spectral slopes.
#' If a character array of dimension names is provided, the spectral slopes
#' will be computed as an average over all elements belonging to those dimensions.
#' If omitted one of c("ftime", "sdate", "time") is searched and the first one with more
#' than one element is chosen.
#' @param verbose Logical for verbose output of the Julia package (default: FALSE).
#' @param drop_realization_dim Logical to remove the "realization" stochastic ensemble dimension (default: FALSE)
#' with the following behaviour if set to TRUE:
#'
#' 1) if \code{nens==1}: the dimension is dropped;
#'
#' 2) if \code{nens>1} and a "member" dimension exists:
#'    the "realization" and "member" dimensions are compacted (multiplied) and the resulting dimension is named "member";
#'
#' 3) if \code{nens>1} and a "member" dimension does not exist: the "realization" dimension is renamed to "member".
#'
#' @return CST_RainFARM() returns a downscaled CSTools object (i.e., of the 
#' class 's2dv_cube').
#' If \code{nens>1} an additional dimension named "realization" is added to the 
#' \code{$data} array after the "member" dimension (unless 
#' \code{drop_realization_dim=TRUE} is specified).
#' The ordering of the remaining dimensions in the \code{$data} element of the input object is maintained.
#' @import multiApply
#' @importFrom utils installed.packages
#' @export
#' @examples
#' \donttest{
#' #Example using CST_RainFARM for a CSTools object
#' nf <- 8   # Choose a downscaling by factor 8
#' exp <- 1 : (2 * 3 * 4 * 8 * 8)
#' dim(exp) <- c(dataset = 1, member = 2, sdate = 3, ftime = 4, lat = 8, lon = 8)
#' lon <- seq(10, 13.5, 0.5)
#' dim(lon) <- c(lon = length(lon))
#' lat <- seq(40, 43.5, 0.5)
#' dim(lat) <- c(lat = length(lat))
#' data <- list(data = exp, lon = lon, lat = lat)
#' # Create a test array of weights
#' ww <- array(1., dim = c(8 * nf, 8 * nf))
#' res <- CST_RainFARM(data, nf, ww, nens=3)
#' str(res)
#' #List of 3
#' # $ data: num [1:4, 1:64, 1:64, 1, 1:2, 1:3] 201 115 119 307 146 ...
#' # $ lon : num [1:64] 9.78 9.84 9.91 9.97 10.03 ...
#' # $ lat : num [1:64] 39.8 39.8 39.9 40 40 ...
#' dim(res$data)
#' # dataset  member   sdate   ftime     lat     lon    realization
#' #       1       2       3       4      64      64     3
#' }
CST_RainFARM <- function(data, nf, weights = 1., slope = 0, kmin = 1,
                         nens = 1, fglob = FALSE, fsmooth = TRUE,
                         time_dim = NULL, verbose = FALSE,
                         drop_realization_dim = FALSE) {

  res <- RainFARM(data$data, data$lon, data$lat,
                  nf, weights, nens, slope, kmin, fglob, fsmooth,
                  time_dim, lon_dim = "lon", lat_dim = "lat",
                  drop_realization_dim, verbose)

  data$data <- res$data
  data$lon <- res$lon
  data$lat <- res$lat

  return(data)
}

#' @rdname RainFARM
#' @title RainFARM stochastic precipitation downscaling (reduced version)
#' @author Jost von Hardenberg - ISAC-CNR, \email{j.vonhardenberg@isac.cnr.it}
#' @description This function implements the RainFARM stochastic precipitation downscaling method
#' and accepts in input an array with named dims ("lon", "lat")
#' and one or more dimension (such as "ftime", "sdate" or "time")
#' over which to average automatically determined spectral slopes.
#' Adapted for climate downscaling and including orographic correction.
#' References:
#' Terzago, S. et al. (2018). NHESS 18(11), 2825–2840. http://doi.org/10.5194/nhess-18-2825-2018,
#' D'Onofrio et al. (2014), J of Hydrometeorology 15, 830-843; Rebora et. al. (2006), JHM 7, 724.
#' Please notice: The function uses the \code{julia_setup()} function from the JuliaCall library, 
#' which takes a long time only the first time it is called on a machine.
#' @param data Precipitation array to downscale.
#' The input array is expected to have at least two dimensions named "lon" and "lat" by default
#' (these default names can be changed with the \code{lon_dim} and \code{lat_dim} parameters)
#' and one or more dimensions over which to average these slopes,
#' which can be specified by parameter \code{time_dim}.
#' The number of longitudes and latitudes in the input data is expected to be even and the same. If not
#' the function will perform a subsetting to ensure this condition.
#' @param lon Vector or array of longitudes.
#' @param lat Vector or array of latitudes.
#' @param weights Matrix with climatological weights which can be obtained using
#' the \code{RFWeights} function. If \code{weights=1.} (default) no weights are used.
#' The matrix should have dimensions (lon, lat) in this order.
#' The names of these dimensions are not checked. 
#' @param nf Refinement factor for downscaling (the output resolution is increased by this factor).
#' @param slope Prescribed spectral slope. The default is \code{slope=0.}
#' meaning that the slope is determined automatically over the dimensions specified by \code{time_dim}.
#' @param kmin First wavenumber for spectral slope (default: \code{kmin=1}).
#' @param nens Number of ensemble members to produce (default: \code{nens=1}).
#' @param fglob Logical to conseve global precipitation over the domain (default: FALSE)
#' @param fsmooth Logical to conserve precipitation with a smoothing kernel (default: TRUE)
#' @param time_dim String or character array with name(s) of time dimension(s)
#' (e.g. "ftime", "sdate", "time" ...) over which to compute spectral slopes.
#' If a character array of dimension names is provided, the spectral slopes
#' will be computed over all elements belonging to those dimensions.
#' If omitted one of c("ftime", "sdate", "time")
#' is searched and the first one with more than one element is chosen.
#' @param lon_dim Name of lon dimension ("lon" by default).
#' @param lat_dim Name of lat dimension ("lat" by default).
#' @param verbose logical for verbose output of the Julia package (default: FALSE).
#' @param drop_realization_dim Logical to remove the "realization" stochastic ensemble dimension (default: FALSE)
#  with the following behaviour if set to TRUE:
#'
#' 1) if \code{nens==1}: the dimension is dropped;
#'
#' 2) if \code{nens>1} and a "member" dimension exists:
#'    the "realization" and "member" dimensions are compacted (multiplied) and the resulting dimension is named "member";
#'
#' 3) if \code{nens>1} and a "member" dimension does not exist: the "realization" dimension is renamed to "member".
#'
#' @return RainFARM() returns a list containing the fine-scale longitudes, latitudes
#' and the sequence of \code{nens} downscaled fields.
#' If \code{nens>1} an additional dimension named "realization" is added to the output array
#' after the "member" dimension (if it exists and unless \code{drop_realization_dim=TRUE} is specified).
#' The ordering of the remaining dimensions in the \code{exp} element of the input object is maintained.
#' @import JuliaCall
#' @import multiApply
#' @export
#' @examples
#' \donttest{
#' # Example for the 'reduced' RainFARM function 
#' nf <- 8   # Choose a downscaling by factor 8
#' nens <- 3 # Number of ensemble members
#' # create a test array with dimension 8x8 and 20 timesteps
#' # or provide your own read from a netcdf file
#' pr <- rnorm(8 * 8 * 20)
#' dim(pr) <- c(lon = 8, lat = 8, ftime = 20)
#' lon_mat <- seq(10, 13.5, 0.5) # could also be a 2d matrix
#' lat_mat <- seq(40, 43.5, 0.5)
#' # Create a test array of weights
#' ww <- array(1., dim = c(8 * nf, 8 * nf))
#' # ... or create proper weights using an external fine-scale climatology file
#' #     Specify a weightsfn filename if you wish to save the weights
#' ww <- RFWeights("./worldclim.nc", nf, lon = lon_mat, lat = lat_mat, 
#'                 weightsfn = "", fsmooth = TRUE)
#' # downscale using weights (ww=1. means do not use weights)
#' res <- RainFARM(pr, lon_mat, lat_mat, nf, 
#'                 fsmooth = TRUE, fglob = FALSE, 
#'                 weights = ww, nens = 2, verbose = TRUE)
#' str(res)
#' #List of 3
#' # $ data: num [1:3, 1:20, 1:64, 1:64] 0.186 0.212 0.138 3.748 0.679 ...
#' # $ lon : num [1:64] 9.78 9.84 9.91 9.97 10.03 ...
#' # $ lat : num [1:64] 39.8 39.8 39.9 40 40 …
#' dim(res$data)
#' #  lon         lat       ftime realization 
#' #   64          64          20           2 
#' }
#'
RainFARM <- function(data, lon, lat, nf, weights = 1., nens = 1,
                     slope = 0, kmin = 1, fglob = FALSE, fsmooth = TRUE,
                     time_dim = NULL, lon_dim = "lon", lat_dim = "lat",
                     drop_realization_dim = FALSE, verbose = FALSE) {

  # Ensure input  grid is square and with even dimensions
  if ( (dim(data)[lon_dim] != dim(data)[lat_dim]) |
       (dim(data)[lon_dim] %% 2 == 1)) {
    print("Warning: input data are expected to be on a square grid")
    print("         with an even number of pixels per side")
    nmin <- min(dim(data)[lon_dim], dim(data)[lat_dim])
    nmin <- floor(nmin / 2) * 2
    data <- .subset(data, lat_dim, 1:nmin)
    data <- .subset(data, lon_dim, 1:nmin)
    if (length(dim(lon)) == 2) {
      lon <- lon[1:nmin, 1:nmin]
      lat <- lat[1:nmin, 1:nmin]
    } else {
      lon <- lon[1:nmin]
      lat <- lat[1:nmin]
    }
    print("The input data have been cut to the range")
    print(paste0("lon: [", lon[1], ", ", lon[length(lon)], "] ",
                 " lat: [", lat[1], ", ", lat[length(lat)], "]"))
  }
  if (!(length(dim(weights)) == 0) &
      !(all(dim(weights) == c(dim(data)[lon_dim] * nf,
                              dim(data)[lat_dim] * nf)))) {
    stop(paste("The dimensions of the weights matrix (", dim(weights)[1],
               "x", dim(weights)[2] ,
               ") are not consistent with the size of the data (",
               dim(data)[lon_dim], ") and the refinement factor (", nf, ")"))
  }
  # Check/detect time_dim
  if (is.null(time_dim)) {
    time_dim_names <- c("ftime", "sdate", "time")
    time_dim_num <- which(time_dim_names %in% names(dim(data)))
    if (length(time_dim_num) > 0) {
      # Find time dimension with length > 1
      ilong <- which(dim(data)[time_dim_names[time_dim_num]] > 1)
      if (length(ilong) > 0) {
        time_dim <- time_dim_names[time_dim_num[ilong[1]]]
      } else {
        stop("No time dimension longer than one found.")
      }
    } else {
      stop("Could not automatically detect a target time dimension ",
           "in the provided data in 'data'.")
    }
    print(paste("Selected time dim:", time_dim))
  }

  # Check availability of JuliaCall
  if (!("JuliaCall" %in% rownames(installed.packages()))) {
    stop("Package 'JuliaCall' needs to be installed in order ",
         "to run RainFARM.")
  }

  # Perform common calls
  JuliaCall::julia_setup(useRCall = FALSE)
  julia_library("RainFARM")

  r <- julia_call("lon_lat_fine", lon, lat, as.integer(nf),
                  need_return = "R")
  lon_f <- r[[1]]
  lat_f <- r[[2]]

  # reorder and group time_dim together at the end
  cdim0 <- dim(data)
  imask <- names(cdim0) %in% time_dim
  data <- aperm(data, c(which(!imask), which(imask)))
  cdim <- dim(data)
  ind <- 1:length(which(!imask))
  # compact (multiply) time_dim dimensions
  dim(data) <- c(cdim[ind], rainfarm_samples = prod(cdim[-ind]))

  # Repeatedly apply .RainFARM
  result <- Apply(data, c(lon_dim, lat_dim, "rainfarm_samples"), .RainFARM,
                  weights, nf, nens, slope, kmin,
                  fglob, fsmooth, verbose)$output1
  # result has dims: lon, lat, rainfarm_samples, realization, other dims
  # Expand back rainfarm_samples to compacted dims
  dim(result) <- c(dim(result)[1:2], cdim[-ind], dim(result)[-(1:3)])

  # Reorder as it was in original data
  # + realization dim after member if it exists
  ienspos <- which(names(cdim0) == "member")
  if ( length(ienspos) == 0) ienspos <- length(names(cdim0))
  iorder <- sapply(c(names(cdim0)[1:ienspos], "realization",
                     names(cdim0)[-(1:ienspos)]),
                   grep, names(dim(result)))
  result <- aperm(result, iorder)

  if (drop_realization_dim) {
    cdim <- dim(result)
    if (nens == 1) {
      # just drop it if only one member
      dim(result) <- cdim[-which(names(cdim) == "realization")[1]]
    } else if ("member" %in% names(cdim)) {
      # compact member and realization dimension if member dim exists,
      # else rename realization to member
      ind <- which(names(cdim) %in% c("member", "realization"))
      dim(result) <- c(cdim[1:(ind[1] - 1)], cdim[ind[1]] * cdim[ind[2]],
                       cdim[(ind[2] + 1):length(cdim)])
    } else {
      ind <- which(names(cdim) %in% "realization")
      names(dim(result))[ind] <- "member"
    }
  }

  return(list(data = result, lon = lon_f, lat = lat_f))
}

#' Atomic RainFARM
#' @param pr Precipitation array to downscale with dimensions (lon, lat, time).
#' @param weights Matrix with climatological weights which can be obtained using
#' the \code{RFWeights} function (default: \code{weights=1.} i.e. no weights).
#' @param nf Refinement factor for downscaling (the output resolution is increased by this factor).
#' @param slope Prescribed spectral slope (default: \code{slope=0.}
#' meaning that the slope is determined automatically over the dimensions specified by \code{time_dim}.
#' @param kmin First wavenumber for spectral slope (default: \code{kmin=1}).
#' @param nens Number of ensemble members to produce (default: \code{nens=1}).
#' @param fglob Logical to conseve global precipitation over the domain (default: FALSE).
#' @param fsmooth Logical to conserve precipitation with a smnoothing kernel (default: TRUE).
#' @param verbose Logical for verbose output of the Julia package (default: FALSE).
#' @return .RainFARM returns a downscaled array with dimensions (lon, lat, time, realization)
#' @noRd
.RainFARM <- function(pr, weights, nf, nens, slope, kmin,
                      fglob, fsmooth, verbose) {

  if (slope == 0) {
    fxp <- julia_call("fft3d", pr, need_return = "R")[[1]]
    sx <- julia_call("fitslopex", fxp, kmin = as.integer(kmin),
                     need_return = "R")
  } else {
    sx <- slope
  }

  result_dims <- c(dim(pr)[1] * nf, dim(pr)[2] * nf, dim(pr)[3],
                   realization = nens)
  r <- array(dim = result_dims)

  for (i in 1:nens) {
    r[, , , i] <- julia_call("rainfarm", pr, sx, as.integer(nf), weights,
                             fglob = fglob, fsmooth = fsmooth,
                             verbose = verbose, need_return = "R")
  }
  return(r)
}

# Function to generalize through do.call() n-dimensional array subsetting 
# and array indexing. Derived from Stack Overflow issue
# https://stackoverflow.com/questions/14500707/select-along-one-of-n-dimensions-in-array
.subset <- function(field, dim_name, range, drop = FALSE) {

  idim  <- which( names(dim(field)) %in% dim_name )
  # Create list representing arguments supplied to [
  # bquote() creates an object corresponding to a missing argument
  indices <- rep(list(bquote()), length(dim(field)))
  indices[[idim]] <- range

  # do.call on the indices
  field <- do.call("[",c(list(field), indices, list(drop = drop)))

  return(field)
}
