#' Make categorical forecast based on a multi-model forecast with different combination methods including multi-model weighting
#'
#'@author Bert Van Schaeybroeck, \email{bertvs@meteo.be}
#'@description This function converts an multi-model ensemble forecast into a categorical forecast. The categories are quantiles (e.g. terciles) and different methods are included to combine the different models. These include ensemble pooling (equal weight for all ensemble members), model combination (equal weight for each model) and model weighting. The weighting method is described in Rajagopalan et al. (2002), Robertson et al. 2004 and Van Schaeybroeck and Vannitsem (2019). The method is based on optimization of the ignorance score and weighs the occurence probabilities of the different models and the climatology.
#'
#'@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}.
#'@param amt.cat is the amount of categories. Equally-sized quantiles will be calculated based on the amount of categories.
#'@param method method used to produce the categorical forecast, can be either "pool", "comb", "mmw" or "obs". The method pool assumes equal weight for all ensemble members while the method comb assumes equal weight for each model. The weighting method is descirbed in Rajagopalan et al. (2002), Robertson et al. 2004 and Van Schaeybroeck and Vannitsem (2019). The obs method returns the categories of the observation. The models corresponding to each ensemble member should be introduced through dimnames of the exp object. All results are obtained through cross validation.
#'
#'@return an object of class \code{s2dv_cube} containing the categorical forecasts in the element called \code{$data}.
#'
#'@references Rajagopalan, B., Lall, U., & Zebiak, S. E. (2002). Categorical climate forecasts through regularization and optimal combination of multiple GCM ensembles. Monthly Weather Review, 130(7), 1792-1811.
#'@references Robertson, A. W., Lall, U., Zebiak, S. E., & Goddard, L. (2004). Improved combination of multiple atmospheric GCM ensembles for seasonal prediction. Monthly Weather Review, 132(12), 2732-2744.
#'@references Van Schaeybroeck, B., & Vannitsem, S. (2019). Postprocessing of Long-Range Forecasts. In Statistical Postprocessing of Ensemble Forecasts (pp. 267-290).
#'
#'@import s2dverification
#'@examples
#'
#'# Example 1
#'mod1 <- 1 : (1 * 3 * 4 * 5 * 6 * 7)
#'dim(mod1) <- c(dataset = 1, member = 3, sdate = 4, ftime = 5, lat = 6, lon = 7)
#'dimnames(mod1)[[2]] <- c("MF", "MF", "UKMO")
#'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_cat_fc(exp = exp, obs = obs, amt.cat = 3, method = "mmw")
#'
#'# Example 2
#'amt.dataset <- 1
#'lst.dataset <- paste0("dataset_", seq(1,amt.dataset))
#'lst.model <- c("MF", "ECMWF", "UKMO")
#'amt.model <- length(lst.model)
#'amt.member.per.model <- c(2, 5, 7)
#'amt.member <- sum(amt.member.per.model)
#'lst.member <- unlist(mapply(function(x, y){rep(x, y)}, lst.model, amt.member.per.model, USE.NAMES = F))
#'lst.member.cat <- unlist(mapply(function(x, y){rep(x, y)}, seq(1,amt.model), amt.member.per.model, USE.NAMES = F))
#'msk.mdl <- sapply(lst.model,function(x){lst.member == x})
#'amt.sdate <- 100
#'lst.sdate <- paste0(1980 + seq(1, amt.sdate), "-01-01")
#'amt.ftime <- 3
#'lst.ftime <- paste0("month_", seq(1, amt.ftime))
#'amt.lon <- 2
#'lst.lon <- paste0("lon_",seq(1,amt.lon))
#'amt.lat <- 3 
#'lst.lat <- paste0("lat_",seq(1,amt.lat))
#'bias <- c(0.3, 0.6, 0.3)
#'slope <- c(1, 1.4, 1.8)
#'#'noise.level.fc <- 0.1 * slope
#'#'noise.level.obs <- 0.1
#'amt.data.pts <- amt.dataset * amt.sdate * amt.ftime * amt.lon * amt.lat
#'obs1 <- array(data = rnorm(amt.data.pts),
#'              dim = c(dataset = amt.dataset, member = 1, sdate = amt.sdate, ftime = amt.ftime, lon = amt.lon, lat = amt.lat),
#'              dimnames = list(dataset = lst.dataset, member = c("obs"), sdate = lst.sdate, ftime = lst.ftime, lon = lst.lon, lat = lst.lat))

#'tmp.noise.obs <- array(data = rnorm(amt.data.pts, 0, noise.level.obs),
#'                       dim = c(dataset = amt.dataset, member = 1, sdate = amt.sdate, ftime = amt.ftime, lon = amt.lon, lat = amt.lat),
#'                       dimnames = list(dataset = lst.dataset, member = c("obs"), sdate = lst.sdate, ftime = lst.ftime, lon = lst.lon, lat = lst.lat))

#'mod1 <- array(data = rnorm(amt.data.pts * amt.member, 0, noise.level.fc),
#'              dim = c(dataset = amt.dataset, member = amt.member, sdate = amt.sdate, ftime = amt.ftime, lon = amt.lon, lat = amt.lat),
#'              dimnames = list(dataset = lst.dataset, member = lst.member, sdate = lst.sdate, ftime = lst.ftime, lon = lst.lon, lat = lst.lat))


#'for(i.mbr in seq(1, amt.member)){
#'  mdl.to.use <- lst.member.cat[i.mbr]
#'  mod1[ , i.mbr, , , , ] <- (obs1[ , 1, , , , ] + mod1[ , i.mbr, , , , ]) / slope[mdl.to.use] + bias[mdl.to.use]
#'}
#'obs1 <- obs1 + tmp.noise.obs 

#'lon <- seq(0, 30, amt.lon)
#'lat <- seq(0, 25, amt.lat)
#'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'

#'amt.cat <- 3
#'cat.comb <- CST_cat_fc(exp, obs, amt.cat, method = "comb")
#'cat.pool <- CST_cat_fc(exp, obs, amt.cat, method = "pool")
#'cat.obs <- CST_cat_fc(exp, obs, amt.cat, method = "obs")
#'cat.mmw <- CST_cat_fc(exp, obs, amt.cat, method = "mmw")
#'rps.comb <- mean((cat.comb$data - cat.obs$data)^2, na.rm = T)
#'rps.pool <- mean((cat.pool$data - cat.obs$data)^2, na.rm = T)
#'rps.mmw <- mean((cat.mmw$data - cat.obs$data)^2, na.rm = T)
#'rps.clim <- mean((1/amt.cat - cat.obs$data)^2, na.rm = T)
#'cat("RPSS comb / pool / mmw: ",1 - rps.comb / rps.clim, 1 - rps.pool / rps.clim, 1 - rps.mmw / rps.clim)
#'
#'str(a)
#'@export

library(abind)


calc.percentiles <- function(data, amt.cat){
  if(amt.cat < 2){
    stop("amount of categories too low: ", amt.cat)
  }
  frac.to.use <- 1/amt.cat
  perc.out <- seq(from = frac.to.use, to = 1 - frac.to.use, by = frac.to.use)
  return(perc.out)
}

calc.quantiles <- function(data, amt.cat){
  perc.to.use <- calc.percentiles(data, amt.cat)
  
  quant.out <- quantile(data, perc.to.use)
  
  return(quant.out)
}

calc.cat <- function(data, quant){
  quant <- c(-Inf,quant,Inf)
  categ.all <- cut(as.vector(data), breaks = quant, labels = F)
  if(!is.null(dim(data))){
    dim(categ.all) <- dim(data)
  }
  return(categ.all)
}


get.mdl.features <- function(mdl.names){
  amt.mbr <- length(mdl.names)
  mdl.diff.names <- unique(mdl.names)
  amt.mdl <- length(mdl.diff.names)
  mdl.msk <- array(F,c(amt.mdl, amt.mbr))
  amt.mbr.per.mdl <- array(0, c(amt.mdl))
  if(amt.mdl == 1 & amt.mbr == 1) {
    mdl.msk = array(T, c(1, 1))
    amt.mbr.per.mdl <- array(1, c(1, 1))
  } else {
    mdl.msk <- t(sapply(mdl.diff.names, function(x){mdl.names == x}))
    amt.mbr.per.mdl <- apply(mdl.msk, c(1), sum,na.rm = T)
  }
  return(list(amt.mbr = amt.mbr, amt.mdl = amt.mdl, 
    amt.mbr.per.mdl = amt.mbr.per.mdl, mdl.msk = mdl.msk))
}


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))
  }
  if(dim != amt.dims + 1){
    amt.dims.out <- amt.dims + 1
    dims.tmp <- seq(1, amt.dims.out)
    dims.tmp[seq(dim, amt.dims.out)] <- c(amt.dims.out, seq(dim,amt.dims.out-1))
    arr.out <- aperm(arr.out, dims.tmp)
  }
  return(arr.out)
}

CST_cat_fc <- function(exp, obs, method = "pool", amt.cat = 3, ...) {
  if (!inherits(exp, "s2dv_cube") || !inherits(exp, "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.")
  }
  
  exp$data <- cat_fc(fc = exp$data, obs = obs$data, method = method, amt.cat = amt.cat, ...)
  exp$Datasets <- c(exp$Datasets, obs$Datasets)
  exp$source_files <- c(exp$source_files, obs$source_files)
  return(exp)
}


cat_fc <- function (fc, obs, method, amt.cat, ...) {
  
  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.")
  }
  
  amt.sdate=dim(fc)["sdate"]
  target.dims <- c("member", "sdate")

  return.feat <- list(amt.cat = amt.cat)
  return.feat$dim <- c(amt.cat, amt.sdate)
  return.feat$name <- c("category", "sdate")
  return.feat$dim.name <- list(paste0("cat_", seq(1, amt.cat)),dimnames(fc)[["sdate"]])
  
  cat_fc_out <- apply.obs.fc(obs = obs,
    fc = fc,
    target.dims = target.dims,
    FUN = .cat_fc, 
    return.feat = return.feat,
    method = method,
    amt.cat = amt.cat,
    ...)
  return(cat_fc_out)
}


apply.obs.fc <- function(obs, fc, target.dims, FUN, return.feat, method, amt.cat, ...){
  dimnames.tmp <- dimnames(fc)
  fc.dims.tmp <- dim(fc)
  dims.out.tmp <- return.feat$dim
  
  obs.fc <- combine.obs.fc(obs, fc)
  names.dim <- names(dim(obs.fc))
  amt.dims <- length(names.dim)
  margin.all <- seq(1, amt.dims)
  matched.dims <- match(target.dims, names.dim)
  margin.to.use <- margin.all[-matched.dims]
  arr.out <- apply(X = obs.fc,
    MARGIN = margin.to.use,
    FUN = FUN, 
    method = method,
    amt.cat = amt.cat,
    ...)
  dims.tmp <- dim(arr.out)
  names.dims.tmp <- names(dim(arr.out))
  if(prod(return.feat$dim) != dims.tmp[1]){
    stop("apply.obs.fc: returned dimensions not as expected: ", prod(return.feat$dim), " and ", dims.tmp[1])
  }
  dim(arr.out) <- c(dims.out.tmp, dims.tmp[-1])
  names(dim(arr.out)) <- c(return.feat$name, names.dims.tmp[c(-1)])
  names.dim[matched.dims] <- return.feat$name
  pos <- match(names.dim, names(dim(arr.out)))
  pos_inv <- match(names(dim(arr.out)), names.dim)
  arr.out <- aperm(arr.out, pos)
  for (i.item in seq(1,length(return.feat$name))){
    dimnames.tmp[[pos_inv[i.item]]] <- return.feat$dim.name[[i.item]]
  }
  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)
}

.cat_fc <- function(obs.fc, amt.cat, method) {
  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)
  dim(var.fc) <- c(amt.mbr, amt.sdate)
  dims.fc <- dim(var.fc)
  
  mdl.dimnames <- dimnames(obs.fc)[["member"]][-1]
  mdl.feat <- get.mdl.features(mdl.dimnames)
  amt.mdl <- mdl.feat$amt.mdl
  amt.coeff <- amt.mdl + 1
  var.cat.fc <- array(NA, c(amt.cat, amt.sdate))
  
  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, drop = F]
    obs.tr <- var.obs[-i.fc, drop = F] 
    obs.ev <- var.obs[i.fc, drop = F]
    amt.sdate.tr <- dim(fc.tr)[2]
    amt.sdate.ev <- dim(fc.ev)[2]
    
    quant.to.use <- calc.quantiles(obs.tr, amt.cat)
    cat.fc.ev <- calc.cat(fc.ev, quant.to.use)
    cat.obs.ev <- calc.cat(obs.ev, quant.to.use)
    
    if(method == "mmw"){
    
      cat.fc.tr <- calc.cat(fc.tr, quant.to.use)
      cat.obs.tr <- calc.cat(obs.tr, quant.to.use)
      
      freq.per.mdl.at.obs <- calc.freq.per.mdl.at.obs(cat.obs = cat.obs.tr, cat.fc = cat.fc.tr,
        mdl.feat = mdl.feat, amt.cat = amt.cat)
      freq.per.mdl.ev <- calc.freq.per.mdl(cat.fc = cat.fc.ev, mdl.feat = mdl.feat, amt.cat = amt.cat)
      
      if(i.fc == 1){
        init.par <- c(rep(1 / amt.coeff, amt.coeff)) * 0.999
      }
      #calculate weights on training dataset
      constr.mtrx <- array(0, c(amt.coeff + 1, amt.coeff))
      for (i.coeff in seq(1, amt.coeff)){
        constr.mtrx[i.coeff, i.coeff] <- 1.
      }
      constr.mtrx[amt.coeff + 1, ] <- -1.
      constr.vec <- c(rep(0., amt.coeff), -1.)
      optim.tmp <- constrOptim(theta = init.par, f = funct.optim, grad = funct.optim.grad, 
                          ui = constr.mtrx, ci = constr.vec,
                          freq.per.mdl.at.obs = freq.per.mdl.at.obs) 
      init.par <- optim.tmp$par * (1 - abs(rnorm(amt.coeff, 0, 0.01)))
          
      var.cat.fc[ , i.fc] <- apply(spr(spr(optim.tmp$par, amt.cat, dim = 2) , amt.sdate.ev, dim = 3) *
        freq.per.mdl.ev[ , , , drop = F], c(2,3), sum, na.rm = T)
    } else if (method == "comb") {
      freq.per.mdl.ev <- calc.freq.per.mdl(cat.fc = cat.fc.ev, mdl.feat = mdl.feat, amt.cat = amt.cat)
      var.cat.fc[ , i.fc] <- apply(freq.per.mdl.ev[-amt.coeff, , , drop = F], c(2, 3), mean, na.rm = T)
    } else if (method == "pool") {
      freq.per.mdl.ev <- calc.freq.per.mdl(cat.fc = cat.fc.ev, mdl.feat = NULL, amt.cat = amt.cat)
      var.cat.fc[ , i.fc] <- freq.per.mdl.ev[1, , ]
    } else if (method == "obs") {
      dim(cat.obs.ev) <- c(1,length(cat.obs.ev))
      freq.per.mdl.ev <- calc.freq.per.mdl(cat.fc = cat.obs.ev, mdl.feat = NULL, amt.cat = amt.cat)
      var.cat.fc[ , i.fc] <- freq.per.mdl.ev[1, , ]
    }
  }
  names(dim(var.cat.fc)) <- c("member", "sdate")
  return(var.cat.fc)
}

funct.optim <- function(par, freq.per.mdl.at.obs){
  return(-mean(log(drop(par %*% freq.per.mdl.at.obs)), na.rm = T))
}

funct.optim.grad <- function(par, freq.per.mdl.at.obs){
  amt.model <- dim(freq.per.mdl.at.obs)[1]
  return(-apply(freq.per.mdl.at.obs/spr(drop(par %*% freq.per.mdl.at.obs), amt.model), c(1), mean, na.rm = T))
}

calc.freq.per.mdl.at.obs <- function(cat.obs, cat.fc, amt.cat, mdl.feat){
  amt.mbr <- dim(cat.fc)[1]
  amt.sdate <- dim(cat.fc)[2]
  amt.mdl <- mdl.feat$amt.mdl
  mdl.msk.tmp <- mdl.feat$mdl.msk
  amt.coeff <- amt.mdl + 1
  msk.fc.obs <- (cat.fc == spr(cat.obs, amt.mbr))
  freq.per.mdl.at.obs <- array(NA, c(amt.coeff, amt.sdate))
  for (i.mdl in seq(1, amt.mdl)){
    freq.per.mdl.at.obs[i.mdl, ] <- apply(msk.fc.obs[mdl.msk.tmp[i.mdl, ], , drop = F],
      c(2), mean, na.rm = T)
  }
  freq.per.mdl.at.obs[amt.coeff, ] = 1 / amt.cat
  return(freq.per.mdl.at.obs)
}

calc.freq.per.mdl <- function(cat.fc, amt.cat, mdl.feat){
  amt.mbr <- dim(cat.fc)[1]
  amt.sdate <- dim(cat.fc)[2]
  if(is.null(mdl.feat)){
    amt.mdl <- 1
    mdl.msk.tmp <- array(T, c(1, amt.mbr))
  } else {
    amt.mdl <- mdl.feat$amt.mdl
    mdl.msk.tmp <- mdl.feat$mdl.msk
  }
  amt.coeff <- amt.mdl + 1
  freq.per.mdl <- array(NA, c(amt.coeff, amt.cat, amt.sdate))
  for (i.mdl in seq(1, amt.mdl)){
    ens.mdl.msk <- mdl.msk.tmp[i.mdl, , drop = F]
    for (i.cat in seq(1, amt.cat)){
      freq.per.mdl[i.mdl, i.cat, ] <- apply(i.cat == cat.fc[ens.mdl.msk, , drop = F],
        c(2), mean, na.rm = T)
    }
  }
  freq.per.mdl[amt.coeff, , ] =  1 / amt.cat
  return(freq.per.mdl)
}
