
library(multiApply)
library(easyVerification)

RPS <- function(exp, obs, n_categories, indices_for_clim, Fair = FALSE, time_dim = 'year', member_dim = 'member', ncores = NULL){
  
  ## Checkings
  if (!is.array(exp)){stop('"exp" must be an array.')}
  if (!is.array(obs)){stop('"obs" must be an array.')}
  if (!identical(as.numeric(dim(exp)[time_dim]), as.numeric(dim(obs)[time_dim]))){stop('"time_dim" must have the same length in "exp" and "obs".')}
  if (!is.numeric(n_categories) | n_categories < 2){stop('"n_categories" must be an integer greater than 1.')}
  if (!is.logical(Fair)){stop('"Fair" must be TRUE or FALSE.')}
  if (!is.character(time_dim)){stop('"time_dim" must be a string.')}
  if (!is.character(member_dim)){stop('"member_dim" must be a string.')}
  if (!is.null(ncores)){
    if (!is.numeric(ncores) | ncores < 1){stop('"ncores" must be either NULL or a positive integer.')}
  }
  
  ## RPS
  rps <- multiApply::Apply(data = list(exp = exp, obs = obs), 
                           target_dims = list(exp = c(time_dim, member_dim), obs = time_dim),
                           n_categories = n_categories, indices_for_clim = indices_for_clim,
                           Fair = Fair, time_dim = time_dim, member_dim = member_dim,
                           fun = .RPS, ncores = ncores)$output1
  return(rps)
}

.RPS <- function(exp, obs, n_categories, indices_for_clim, Fair = FALSE, time_dim = 'year', member_dim = 'member'){
  
  exp_probs <- .get_probs(data = exp, type = 'exp', indices_for_quartiles = indices_for_clim, 
                          time_dim = time_dim, member_dim = member_dim, n_categories = n_categories)$probs
  
  obs_probs <- .get_probs(data = obs, type = 'obs', indices_for_quartiles = indices_for_clim, 
                          time_dim = time_dim, member_dim = NULL, n_categories = n_categories)$probs
  
  rps <- .RPS_from_probs(probs_exp = exp_probs, probs_obs = obs_probs, n_categories = n_categories, Fair = Fair, n_members = as.numeric(dim(exp)[member_dim]))
  
  return(rps)
}

.get_probs <- function(data, type, indices_for_quartiles, time_dim, member_dim, n_categories){
  
  if (type == 'exp'){
    target_dims_quantiles = c(time_dim, member_dim)
    target_dims_probs = list(member_dim, 'bin')
  } else if (type == 'obs'){
    target_dims_quantiles = time_dim
    target_dims_probs = list(NULL,'bin')
  } else {stop("Parameter 'type' must be 'exp' or 'obs'")}
  
  quantiles <- multiApply::Apply(data = ClimProjDiags::Subset(x = data, along = time_dim, indices = indices_for_quartiles), 
                                 target_dims = target_dims_quantiles, 
                                 fun = .quantile_aux, n_categories = n_categories, 
                                 output_dims = 'bin', ncores = NULL)$output1
  
  probs <- multiApply::Apply(data = list(data = data, quantiles = quantiles),
                             target_dims = target_dims_probs,
                             fun = .c2p, 
                             output_dims = 'bin', ncores = NULL)$output1
  
  return(list('probs' = probs, 'quantiles' = quantiles))
}

.quantile_aux = function(data, n_categories){
  data = as.vector(data)
  q = quantile(x = data, probs = 1:(n_categories - 1) / n_categories, type = 8, na.rm = TRUE)
  return(q)
}

.c2p <- function(data, quantiles) {
  if (anyNA(data)){
    output <- array(data = NA, dim = length(quantiles) + 1)
  } else {
    output <- colMeans(easyVerification::convert2prob(as.vector(data), threshold = as.vector(quantiles)))
  }
  return(output)
}


.RPS_from_probs <- function(probs_exp, probs_obs, n_categories, Fair, n_members){
  
  rps <- NULL
  for (i in 1:n_categories){
    rps[i] <- (cumsum(probs_exp)[i]-cumsum(probs_obs)[i])^2 ## (Wilks 2011, pp.348-350)
  }
  rps <- sum(rps)
  
  if (Fair == TRUE){ ## FairRPS
    ## adjustment <- rowSums(-1 * (1/R - 1/R.new) * ens.cum * (R - ens.cum)/R/(R - 1)) [formula taken from SpecsVerification::EnsRps]
    R <- n_members
    R_new <- Inf 
    probs_cum <- cumsum(probs_exp)
    adjustment <- -1 / (R-1) * probs_cum * (1 - probs_cum)
    adjustment <- sum(adjustment)
    rps <- rps + adjustment
  }
  
  return(rps)
}
