# This module load the outputs saved for each individual forecast system
# and creates the multimodel ensemble

source("modules/Loading/R/dates2load.R")
source("modules/Loading/R/get_timeidx.R")
source("modules/Loading/R/check_latlon.R")
source('modules/Multimodel/load_multimodel.R')

Multimodel <- function(recipe) {
  
  # recipe: auto-s2s recipe as provided by read_yaml
  
  # Loading data saved in the jobs for individual models
  data <- load_multimodel(recipe)
  
  # Creating the multi-model
  if (tolower(recipe$Analysis$Datasets$Multimodel$approach) == 'pooled') {
    
    # Deterministic hindcast
    data$hcst$data <- CSTools::MergeDims(data = data$hcst$data, 
                                         merge_dims = c('model','ensemble'), 
                                         rename_dim = 'ensemble', na.rm = TRUE)
    data$hcst$dims <- dim(data$hcst$data)
    
    # Deterministic forecast
    if (!is.null(recipe$Analysis$Time$fcst_year)) {
      data$fcst$data <- CSTools::MergeDims(data = data$fcst$data, 
                                           merge_dims = c('model','ensemble'), 
                                           rename_dim = 'ensemble', na.rm = TRUE)
      data$fcst$dims <- dim(data$fcst$data)
    }
    
    # Probabilistic hindcast and forecast
    prob <- Probabilities(recipe, data)
    
  } else if (tolower(recipe$Analysis$Datasets$Multimodel$approach) %in% c('mean','median')) {
    
    # Probabilistic hindcast and forecast
    warning('Probabilities for multi-model mean is still under development.
             If Skill is used, the results are not correct for the probabilitic metrics.
             Probabilities cannot be computed because data has an extra dimension (model).
             Also, the function should return the observed probabilities.
             Maybe it is better to use GetProbs.')
    prob <- NULL #Probabilities(recipe, data)
    
    # Deterministic hindcast
    data$hcst$data <- multiApply::Apply(data = data$hcst$data, 
                                        target_dims = 'ensemble',
                                        fun = tolower(recipe$Analysis$Datasets$Multimodel$approach),
                                        na.rm = TRUE,
                                        ncores = recipe$Analysis$ncores)$output1
    data$hcst$data <- multiApply::Apply(data = data$hcst$data,
                                        target_dims = 'model',
                                        fun = tolower(recipe$Analysis$Datasets$Multimodel$approach),
                                        na.rm = FALSE,
                                        ncores = recipe$Analysis$ncores)$output1
    data$hcst$dims <- s2dv::InsertDim(data = data$hcst$dims, posdim = 6, lendim = 1, name = 'ensemble')
    data$hcst$dims <- dim(data$hcst$data)
    
    # Deterministic forecast
    if (!is.null(recipe$Analysis$Time$fcst_year)) {
      data$fcst$data <- multiApply::Apply(data = data$fcst$data, 
                                          target_dims = 'ensemble',
                                          fun = tolower(recipe$Analysis$Datasets$Multimodel$approach),
                                          na.rm = TRUE,
                                          ncores = recipe$Analysis$ncores)$output1
      data$fcst$data <- multiApply::Apply(data = data$fcst$data,
                                          target_dims = 'model',
                                          fun = tolower(recipe$Analysis$Datasets$Multimodel$approach),
                                          na.rm = FALSE,
                                          ncores = recipe$Analysis$ncores)$output1
      data$fcst$dims <- s2dv::InsertDim(data = data$fcst$dims, posdim = 6, lendim = 1, name = 'ensemble')
      data$fcst$dims <- dim(data$fcst$data)
    }
    
  } else {stop('Incorrect multi-model approach')}
  
  # Removing the temporary data of models in case user only requests multimodel
  if (tolower(recipe$Analysis$Datasets$Multimodel$execute) %in% c('true','yes')) {
    if (tolower(recipe$Analysis$Horizon) == 'seasonal') {
      exp.name <- recipe$Analysis$Datasets$System$models
    } else if (tolower(recipe$Analysis$Horizon) == 'decadal') {
      exp.name <- sapply(recipe$Analysis$Datasets$System$models, '[[', 'name')
    } else {stop('Multimodel not implemented for this horizon.')}
    unlink(file.path(recipe$Run$output_dir,"outputs/*/*",gsub('\\.','',exp.name)),
           recursive = T)
    unlink(file.path(recipe$Run$output_dir,"plots/*",gsub('\\.','',exp.name)),
           recursive = T)
  }
  
  return(list(data = data, prob = prob))
}
