
#' Ensemble Bias Correction based on adjustment of ensemble mean, and ensemble spread
#'
#'@author Bert Van Schaeybroeck, \email{bertvs@meteo.be}
#'@description This function applies a member-by-member ensemble bias correction described in Van Schaeybroeck and Vannitsem (Quart J, 2015). The adjusted forecasts has an optimized CRPS score. This implies the (near) correspondence of 1) forecast mean with observational mean, 2) forecast variability with observational variability, 3) mean squared error with average ensemble variability such that the ensemble is reliable.
#'
#'@param exp an object of class \code{s2dv_cube} as returned by \code{CST_Load} function, containing the seasonal forecast experiment data in the element named \code{$data}
#'@param obs an object of class \code{s2dv_cube} as returned by \code{CST_Load} function, containing the observed data in the element named \code{$data}.
#'
#'@return an object of class \code{s2dv_cube} containing the bias corrected forecasts in the element called \code{$data} with the same dimensions of the experimental data.
#'
#'@references Van Schaeybroeck, B., & Vannitsem, S. (2015). Ensemble post‐processing using member‐by‐member approaches: theoretical aspects. Quarterly Journal of the Royal Meteorological Society, 141(688), 807-818.
#'
#'@import s2dverification
#'@examples
#'
#'# Example
#'# Creation of sample s2dverification objects. These are not complete
#'# s2dverification objects though. The Load function returns complete objects.
#'mod1 <- 1 : (1 * 3 * 4 * 5 * 6 * 7)
#'dim(mod1) <- c(dataset = 1, member = 3, sdate = 4, ftime = 5, lat = 6, lon = 7)
#'obs1 <- 1 : (1 * 1 * 4 * 5 * 6 * 7)
#'dim(obs1) <- c(dataset = 1, member = 1, sdate = 4, ftime = 5, lat = 6, lon = 7)
#'lon <- seq(0, 30, 5)
#'lat <- seq(0, 25, 5)
#'exp <- list(data = mod1, lat = lat, lon = lon)
#'obs <- list(data = obs1, lat = lat, lon = lon)
#'attr(exp, 'class') <- 's2dv_cube'
#'attr(obs, 'class') <- 's2dv_cube'
#'a <- CST_EnsBiasCorrection(exp = exp, obs = obs)
#'str(a)
#'@export

library(abind)

spr <- function(x, amt.spr, dim = 1) {
  if(is.vector(x)){
    amt.dims <- 1
    arr.out <- array(rep(x, amt.spr), c(length(x), amt.spr))
  } else if(is.array(x)) {
    amt.dims <- length(dim(x))
    arr.out <- array(rep(as.vector(x), amt.spr), c(dim(x), amt.spr))
  } else {
    stop("x is not array nor vector but is ", class(x))
  }
  amt.dims.out <- amt.dims + 1
  dims.tmp <- seq(1, amt.dims.out)
  dims.tmp[seq(dim, amt.dims.out)] <- c(amt.dims + 1, seq(dim,amt.dims.out-1))
  arr.out <- aperm(arr.out, dims.tmp)
  return(arr.out)
}

calc.crps.opt <- function(par, quant.obs.fc){
  return(
    with(quant.obs.fc, 
      mean(
        apply(abs(obs.per.ens - (par[1] + par[2] * fc.ens.av.per.ens +
	    ((par[3])^2 + par[4] / spr.abs.per.ens) * fc.dev)), c(2), mean, na.rm = T) -
        abs((par[3])^2 * spr.abs + par[4]) / 2.
      , na.rm = T)
    )
  )
}

calc.crps <- function(obs, fc){
  quant.obs.fc <- calc.obs.fc.quant(obs = obs, fc = fc)
  return(with(quant.obs.fc,
    mean(apply(abs(obs.per.ens - fc), c(2), mean, na.rm = T) - spr.abs / 2., na.rm = T)))
}

calc.obs.fc.quant <- function(obs, fc){
  amt.mbr <- dim(fc)[1]
  obs.per.ens <- spr(obs, amt.mbr)
  fc.ens.av <- apply(fc, c(2), mean, na.rm = T)
  cor.obs.fc <- cor(fc.ens.av, obs, use = "complete.obs")
  obs.av <- mean(obs, na.rm = T)
  obs.sd <- sd(obs, na.rm = T)
  return(
    append(
      calc.fc.quant(fc = fc),
      list(
        obs.per.ens = obs.per.ens,
        cor.obs.fc = cor.obs.fc,
        obs.av = obs.av,
        obs.sd = obs.sd
      )
    )
  )
}


calc.fc.quant <- function(fc){
  amt.mbr <- dim(fc)[1]
  fc.ens.av <- apply(fc, c(2), mean, na.rm = T)
  fc.ens.av.av <- mean(fc.ens.av, na.rm = T)
  fc.ens.av.per.ens <- spr(fc.ens.av, amt.mbr)
  fc.ens.sd <- apply(fc, c(2), sd, na.rm = T)
  fc.ens.sd.av <- mean(fc.ens.sd,na.rm = T)
  fc.dev <- fc - fc.ens.av.per.ens
  repmat1.tmp <- spr(fc, amt.mbr)
  repmat2.tmp <- aperm(repmat1.tmp, c(2, 1, 3))
  spr.abs <- apply(abs(repmat1.tmp - repmat2.tmp), c(3), mean, na.rm = T)
  spr.abs.per.ens <- spr(spr.abs, amt.mbr)
  fc.av <- mean(fc, na.rm = T)
  fc.sd <- sd(fc, na.rm = T)
  return(
    list(
      fc.ens.av = fc.ens.av,
      fc.ens.av.av = fc.ens.av.av,
      fc.ens.av.per.ens = fc.ens.av.per.ens,
      fc.ens.sd = fc.ens.sd,
      fc.ens.sd.av = fc.ens.sd.av,
      fc.dev = fc.dev,
      spr.abs = spr.abs,
      spr.abs.per.ens = spr.abs.per.ens,
      fc.av = fc.av,
      fc.sd = fc.sd
    )
  )
}

calc.wer.cr.par <- function(quant.obs.fc){
  par.out <- rep(NA, 3)
  par.out[3] <- with(quant.obs.fc, sqrt(obs.sd * sqrt(1 - cor.obs.fc^2) / mean(fc.ens.sd.av, na.rm = T)))
  par.out[2] <- with(quant.obs.fc, cor.obs.fc * obs.sd / fc.sd)
  par.out[1] <- with(quant.obs.fc, obs.av - par.out[2] * fc.ens.av.av, na.rm = T)
  return(par.out)
}

correct.mbm.fc <- function(fc, par){
  quant.fc.mp <- calc.fc.quant(fc = fc)
  return(with(quant.fc.mp, par[1] + par[2] * fc.ens.av.per.ens + fc.dev * abs((par[3])^2 + par[4] / spr.abs)))
}

CST_EnsBiasCorrection <- function(fc, obs) {
  if (!inherits(exp, "s2dv_cube") || !inherits(fc, "s2dv_cube")) {
    stop("Parameter 'exp' and 'obs' must be of the class 's2dv_cube', ",
         "as output by CSTools::CST_Load.")
  }

  if (dim(obs$data)["member"] != 1) {
    stop("The length of the dimension 'member' in the component 'data' ",
         "of the parameter 'obs' must be equal to 1.")
  }

  fc$data <- EnsBiasCorrection(fc = fc$data, obs = obs$data)
  fc$Datasets <- c(fc$Datasets, obs$Datasets)
  fc$source_files <- c(fc$source_files, obs$source_files)
  return(fc)
}

EnsBiasCorrection <- function (fc, obs) {
 
   if (!all(c("member", "sdate") %in% names(dim(fc)))) {
    stop("Parameter 'fc' must have the dimensions 'member' and 'sdate'.")
  }

  if (!all(c("sdate") %in% names(dim(obs)))) {
    stop("Parameter 'obs' must have the dimension 'sdate'.")
  }

  if (any(is.na(fc)))  {
    warning("Parameter 'fc' contains NA values.")
  }
 
  if (any(is.na(obs)))  {
    warning("Parameter 'obs' contains NA values.")
  }

  target.dims <- c("member", "sdate")
  EnsBiasCorrected <- apply.obs.fc(obs = obs,
                      fc = fc,
                      target.dims = target.dims,
                      FUN = .esbc)
                      
  return(EnsBiasCorrected)
}

apply.obs.fc <- function(obs, fc, target.dims, FUN, ...){
  dimnames.tmp <- dimnames(fc)
  fc.dims.tmp <- dim(fc)
  dims.out.tmp <- fc.dims.tmp[target.dims]
  obs.fc <- combine.obs.fc(obs, fc)
  names.dim <- names(dim(obs.fc))
  amt.dims <- length(names.dim)
  margin.all <- seq(1, amt.dims)
  margin.to.use <- margin.all[-match(target.dims, names.dim)]
  arr.out <- apply(X = obs.fc,
               MARGIN = margin.to.use,
               FUN = FUN)        
  dims.tmp <- dim(arr.out)
  names.dims.tmp <- names(dim(arr.out))
  dim(arr.out) <- c(dims.out.tmp, dims.tmp[-1])
  names(dim(arr.out)) <- c(target.dims, names.dims.tmp[c(-1)])
  pos <- match(names.dim, names(dim(arr.out)))
  arr.out <- aperm(arr.out, pos)
  dimnames(arr.out) <- dimnames.tmp
  return(arr.out)
}

combine.obs.fc <- function(obs,fc){
    names.dim.tmp <- names(dim(obs))
    members.dim <- which(names.dim.tmp == "member")
    arr.out <- abind(obs, fc, along = members.dim)
    dimnames.tmp <- dimnames(arr.out)
    names(dim(arr.out)) <- names.dim.tmp
    dimnames(arr.out) <- dimnames.tmp
    names(dimnames(arr.out)) <- names.dim.tmp
    return(arr.out)
}

.esbc <- function(obs.fc) {
  dims.tmp=dim(obs.fc)
  amt.mbr <- dims.tmp["member"][]-1
  amt.sdate <- dims.tmp["sdate"][]
  pos <- match(c("member","sdate"), names(dims.tmp))
  obs.fc <- aperm(obs.fc, pos)
  var.obs <- asub(obs.fc, list(1),1)
  var.fc <- asub(obs.fc, list(1+seq(1, amt.mbr)),1)
  dims.fc <- dim(var.fc)
  var.cor.fc <- NA * var.fc
    
 
  for (i.fc in seq(1, amt.sdate)) {
    # defining forecast,hindcast and observation in cross-validation
    fc.ev <- var.fc[ , i.fc, drop = F]
    fc.tr <- var.fc[ , -i.fc]
    obs.tr <- var.obs[-i.fc, drop = F] 
    
    #calculate ensemble and observational characteristics
    quant.obs.fc.tr <- calc.obs.fc.quant(obs = obs.tr, fc = fc.tr)
    #calculate initial value for regression parameters
    init.par <- c(calc.wer.cr.par(quant.obs.fc.tr), 0.001)
    #calculate regression parameters on training dataset
    optim.tmp <- optim(par = init.par, fn = calc.crps.opt
      , quant.obs.fc = quant.obs.fc.tr) #, control = list(maxit = 2,trace = T)
    mbm.par <- optim.tmp$par
    var.cor.fc[ , i.fc] <- correct.mbm.fc(fc.ev , mbm.par)
  }
  
  cat("CRPS uncorr/init/optim/eval: ",
      calc.crps(var.obs, var.fc),
      calc.crps(var.obs, var.cor.fc),"\n")
  
  names(dim(var.cor.fc)) <- c("member", "sdate")
  return(var.cor.fc)
}



