#'CST_AdamontQQCorr computes quantile-quantile correction of seasonal or decadal forecast data using weather types
#'
#'@description This function computes a quantile mapping based on weather types for experiment data (typically a hindcast) onto reference \code{obs}, typically provided by reanalysis data.
#'@author Paola Marson, \email{paola.marson@meteo.fr} for PROSNOW version
#'@author Lauriane Batté, \email{lauriane.batte@meteo.fr} for CSTools adaptation
#'
#'@param exp experiment data an object of class \code{s2dv_cube}
#'@param wt_exp corresponding weather types (same dimensions as \code{exp$data} but lat/lon)
#'@param obs reference data, also of class \code{s2dv_cube}. lat/lon dimensions can differ from \code{exp} if non rectilinear latlon grids are used, in which case regrid should be set to TRUE and AdamontNearestNeighbors \code{NN} output should be provided
#'@param corrdims list of dimensions in \code{exp} for which quantile mapping correction is applied
#'@param londim character name of longitude dimension in \code{exp} and \code{obs}
#'@param latdim character name of latitude dimension in \code{exp} and \code{obs}
#'@param regrid boolean indicating whether Nearest Neighbors regridding is needed
#'@param NN list (output from Adamont_NearestNeighbors) maps (nlat, nlon) onto (nlat_o, nlon_o)
#'
#'@import qmap
#'@import multiApply
#'@import abind
#'@examples
#'\dontrun{
#'wt_exp <- sample(1:3, 15*6*3, replace=T)
#'dim(wt_exp) <- c(dataset=1, member=15, sdate=6, ftime=3)
#'wt_obs <- sample(1:3, 6*3, replace=T)
#'dim(wt_obs) <- c(dataset=1, member=1, sdate=6, ftime=3)
#'exp_corr <- CST_AdamontQQCorr(exp=lonlat_data$exp, wt_exp=wt_exp, obs=lonlat_data$obs, wt_obs=wt_obs, corrdims = c('dataset','member','sdate','ftime'))
#'}
CST_AdamontQQCorr <- function(exp, wt_exp, obs, wt_obs, corrdims = c('member','sdate','ftime'), londim='lon', latdim='lat') {

   if (!inherits(exp, 's2dv_cube') || !inherits(obs, 's2dv_cube')){
     stop("Inputs 'exp' and 'obs' must be of class 's2dv_cube', ",
	  "as output by CSTools::CST_Load.")
   }
   dimnames <- names(dim(obs$data))
   if (is.null(which(dimnames == latdim)) || is.null(which(dimnames==londim))){
     stop("'londim' or 'latdim' input doesn't match with 'obs' dimension names")   
   }
   if ((length(dim(exp$lon))==2) || (length(dim(obs$lon))==2)){
      myNN <- NearestNeighbors(exp=exp, obs=obs, method='ADA')
      exp_corr <- AdamontQQCorr(exp=exp$data, wt_exp=wt_exp, obs=obs$data, wt_obs=wt_obs, corrdims=corrdims, londim=londim, latdim=latdim, regrid=TRUE, NN=myNN) 
   } else {
   # If not (standard case)
      exp_corr <- AdamontQQCorr(exp=exp$data, wt_exp=wt_exp, obs=obs$data, wt_obs=wt_obs, corrdims=corrdims, londim=londim, latdim=latdim, regrid=FALSE)
   }
   return(exp_corr)
}


#'AdamontQQCorr computes quantile-quantile correction of seasonal or decadal forecast data using weather types
#'
#'@description This function computes a quantile mapping based on weather types for experiment data (typically a hindcast) onto reference \code{obs}, typically provided by reanalysis data.
#'@author Paola Marson, \email{paola.marson@meteo.fr} for PROSNOW version
#'@author Lauriane Batté, \email{lauriane.batte@meteo.fr} for CSTools adaptation
#'
#'@param exp array with named dimensions (such as \code{$data} array of experiment data from an object of class \code{s2dv_cube})
#'@param wt_exp corresponding weather types (same dimensions as \code{exp} but lat/lon)
#'@param obs array with named dimensions with reference data (can also be \code{$data} array of class \code{s2dv_cube}). lat/lon dimensions can differ from \code{exp} if non rectilinear latlon grids are used, in which case regrid should be set to TRUE and AdamontNearestNeighbors \code{NN} output should be provided
#'@param corrdims list of dimensions in \code{exp} for which quantile mapping correction is applied
#'@param londim character name of longitude dimension in \code{exp} and \code{obs}
#'@param latdim character name of latitude dimension in \code{exp} and \code{obs}
#'@param regrid boolean indicating whether Nearest Neighbors regridding is needed
#'@param NN list (output from AdamontNearestNeighbors) maps (nlat, nlon) onto (nlat_o, nlon_o)
#'
#'@import qmap
#'@import multiApply
#'@import abind
#'@examples
#'\dontrun{
#'wt_exp <- sample(1:3, 15*6*3, replace=T)
#'dim(wt_exp) <- c(dataset=1, member=15, sdate=6, ftime=3)
#'wt_obs <- sample(1:3, 6*3, replace=T)
#'dim(wt_obs) <- c(dataset=1, member=1, sdate=6, ftime=3)
#'exp_corr <- AdamontQQCorr(exp=lonlat_data$exp$data, wt_exp=wt_exp, obs=lonlat_data$obs$data, wt_obs=wt_obs, corrdims = c('dataset','member','sdate','ftime'))
#'}
AdamontQQCorr <- function(exp, wt_exp, obs, wt_obs, corrdims = c('member', 'sdate', 'ftime'), londim='lon', latdim='lat', regrid=FALSE, NN=NULL) {

   # The regridding part should only be done if lat/lon dimensions of obs and
   # exp differ.
   # INCLUDE SANITY CHECKS: LATDIM AND LONDIM MUST ALSO MATCH EXP LAT/LON DIM NAMES
   # IF REGRID = TRUE AND NN = NULL: error !!
   if (regrid == 'TRUE'){
      obsdims <- names(dim(obs))
      poslat <- which(obsdims == latdim)
      poslon <- which(obsdims == londim)
      nlat_o <- dim(obs)[poslat]
      nlon_o <- dim(obs)[poslon]
      ilat_o <- array(c(1:nlat_o))
      names(dim(ilat_o))[1] <- latdim
      ilon_o <- array(c(1:nlon_o))
      names(dim(ilon_o))[1] <- londim
      # First step if obs data is higher resolution than exp data is to use nearest neighbor
      # to compute downscaling of exp data	
      exp_corr <- Apply(list(exp,ilat_o,ilon_o),target_dims=list(c(latdim,londim),latdim,londim),.getNN,NN=NN)$output1

      # Reorder exp_corr dimensions to match exp dimensions
      dexpc <- match(names(dim(exp)), names(dim(exp_corr)))
      exp_corr <- aperm(exp_corr,dexpc)
      dimnames(exp_corr) <- dimnames(exp)[dexpc]
      # Keep original wt_exp for remapping data
      wt_exp2 <- wt_exp

      # Both exp and obs data are now on the same grid
   } else {
      exp_corr <- exp
      # Keep original wt_exp for remapping data
      wt_exp2 <- wt_exp
   }
   
   # Use CST_QuantileMapping function for quantile mapping, depending on weather type
   for (i in 1:(length(corrdims) - 1)) {
      obs <- MergeDims(obs, corrdims[i:(i+1)], rename=corrdims[i+1])
      wt_obs <- MergeDims(wt_obs, corrdims[i:(i+1)], rename=corrdims[i+1])
      exp_corr <- MergeDims(exp_corr, corrdims[i:(i+1)], rename=corrdims[i+1])
      wt_exp2 <- MergeDims(wt_exp2, corrdims[i:(i+1)], rename=corrdims[i+1])
   }
   names(dim(obs))[which(names(dim(obs)) == corrdims[length(corrdims)])] <- 'time'
   names(dim(wt_obs))[which(names(dim(wt_obs)) == corrdims[length(corrdims)])] <- 'time'
   names(dim(exp_corr))[which(names(dim(exp_corr)) == corrdims[length(corrdims)])] <- 'time'
   names(dim(wt_exp2))[which(names(dim(wt_exp2)) == corrdims[length(corrdims)])] <- 'time'
   # Split 'time' dim in weather types
   obs <- SplitDim(obs, split_dim='time',indices=as.vector(wt_obs), new_dim_name='type')
   exp_corr <- SplitDim(exp_corr, split_dim='time',indices=as.vector(wt_exp2), new_dim_name='type')
   # Add NAs to exp_corr if needed to have compatible sample dimensions
   numtobs <- dim(obs)[which(names(dim(obs))=='time')]
   numtexp <- dim(exp_corr)[which(names(dim(exp_corr))=='time')]
   if (numtexp%%numtobs > 0){
      # Create extra dimension and include NAs
      ndimexp <- names(dim(exp_corr))
      ndimobs <- names(dim(obs))
      postime <- which(ndimexp=='time')
      dimadd <- dim(exp_corr)
      dimadd[postime] <- ceiling(numtexp/numtobs)*numtobs-numtexp
      exp_corr <- abind::abind(exp_corr,array(NA,dimadd),along=postime)
      names(dim(exp_corr)) <- ndimexp
      exp_corr <- SplitDim(exp_corr,'time',freq=numtobs,indices=NULL)
      dimobs <- c(dim(obs),1)
      dim(obs) <- dimobs
      names(dim(obs)) <- c(ndimobs,'index')
      res <- QuantileMapping(exp=exp_corr,obs=obs,sample_dims=c('time','index'),method='RQUANT')
      res <- MergeDims(res,c('time','index'))
      # Remove the extra NA values added previously
      res <- Subset(res,along='time',indices=1:numtexp)
   } else {
      # Apply QuantileMapping to exp_corr depending on weather type
      res <- QuantileMapping(exp=exp_corr,obs=obs,sample_dims='time',samplemethod='RQUANT')
   }
   rm(exp_corr) # Save space in memory
   # Reshape exp_corr data onto time dimension before 'Split'
   rep_pos <- array(NA,c(time=length(wt_exp2)))
   pos_time <- which(names(dim(res)) == 'time')
   pos_type <- which(names(dim(res)) == 'type')
   for (x in unique(wt_exp2)){
      rep_pos[which(wt_exp2==x)]<-1:length(which(wt_exp2==x))
   }
   exp_corr <- .unsplit_wtype(exp=res,wt_exp=wt_exp2,rep_pos=rep_pos,pos_time=pos_time) 
   # Now reshape exp_corr data onto original dimensions
   dim(exp_corr) <- c(dim(wt_exp), dim(exp_corr)[-c(pos_time,pos_type)])
   return(exp_corr)
}

.getNN <- function(exp,ilat,ilon,NN){
   return(exp[NN$imin_lat[ilat,ilon],NN$imin_lon[ilat,ilon]])
}

.unsplit_wtype <- function(exp=exp,dim_wt='type',wt_exp=wt_exp,dim_time='time',rep_pos=rep_pos,pos_time=1){
   # Initiate output
   new <- Subset(Subset(exp, along=dim_wt, indices=wt_exp[1]), along=dim_time, indices=rep_pos[1])
   dimnames <- names(dim(new))
   for (x in 2:length(wt_exp)){
     	dat <- Subset(Subset(exp, along=dim_wt, indices=wt_exp[x]), along=dim_time, indices=rep_pos[x])
        new <- abind::abind(new,dat,along=pos_time)
   }
   names(dim(new)) <- dimnames
   return(new)
}

