## TODO: Remove in the next release
source("modules/Calibration/calibrate_datasets.R")

Calibration <- function(recipe, data, retrieve = TRUE) {
  # Function that calibrates the hindcast using the method stated in the
  # recipe. If the forecast is not null, it calibrates it as well.
  #
  # data: list of s2dv_cube objects containing the hcst, obs and fcst.
  #       Optionally, it may also have hcst.full_val and obs.full_val.
  # recipe: object obtained when passing the .yml recipe file to read_yaml()
  
  method <- tolower(recipe$Analysis$Workflow$Calibration$method)

  if (method == "raw") {
    warn(recipe$Run$logger, retrieve = retrieve,
         paste("The Calibration module has been called, but the calibration",
               "method in the recipe is 'raw'. The hcst and fcst will not be",
               "calibrated."))
    fcst_calibrated <- data$fcst
    hcst_calibrated <- data$hcst
    if (!is.null(data$hcst.full_val)) {
      hcst_full_calibrated <- data$hcst.full_val
    } else {
      hcst_full_calibrated <- NULL
    }
    CALIB_MSG <- "##### NO CALIBRATION PERFORMED #####"

  } else {
    ## TODO: Calibrate full fields when present
    # Calibration function params
    mm <- !is.null(recipe$Analysis$Datasets$Multimodel) && 
          !tolower(recipe$Analysis$Datasets$Multimodel) %in% c('no','false')
    if (is.null(recipe$Analysis$ncores)) {
      ncores <- 1
    } else {
      ncores <- recipe$Analysis$ncores
    }
    if (is.null(recipe$Analysis$remove_NAs)) {
      na.rm = F
    } else {
      na.rm = recipe$Analysis$remove_NAs
    }
    
    CALIB_MSG <- "##### CALIBRATION COMPLETE #####"
    # Replicate observation array for the multi-model case
    
    if (recipe$Analysis$Variables$freq == "monthly_mean") {
    
      CST_CALIB_METHODS <- c("bias", "evmos", "mse_min", "crps_min", "rpc-based")
      ## TODO: implement other calibration methods
      if (!(method %in% CST_CALIB_METHODS)) {
        error(recipe$Run$logger, retrieve = retrieve, 
              paste("Calibration method in the recipe is not available for",
                    "monthly data."))
        stop()
      } else {
        # Calibrate the hindcast
        hcst_calibrated <- CST_Calibration(data$hcst, data$obs,
                                           cal.method = method,
                                           eval.method = "leave-one-out",
                                           multi.model = mm,
                                           na.fill = TRUE,
                                           na.rm = na.rm,
                                           apply_to = NULL,
                                           alpha = NULL,
                                           memb_dim = "ensemble",
                                           sdate_dim = "syear",
                                           ncores = ncores)
        # In the case where anomalies have been computed, calibrate full values
        if (!is.null(data$hcst.full_val)) {
          hcst_full_calibrated <- CST_Calibration(data$hcst.full_val,
                                                  data$obs.full_val,
                                                  cal.method = method,
                                                  eval.method = "leave-one-out",
                                                  multi.model = mm,
                                                  na.fill = TRUE,
                                                  na.rm = na.rm,
                                                  apply_to = NULL,
                                                  memb_dim = "ensemble",
                                                  sdate_dim = "syear",
                                                  ncores = ncores)
        } else {
          hcst_full_calibrated <- NULL
        }

        # Calibrate the forecast
        if (!is.null(data$fcst)) {
          fcst_calibrated <- CST_Calibration(data$hcst, data$obs, data$fcst,
                                             cal.method = method,
                                             eval.method = "leave-one-out",
                                             multi.model = mm,
                                             na.fill = TRUE,
                                             na.rm = na.rm,
                                             apply_to = NULL,
                                             alpha = NULL,
                                             memb_dim = "ensemble",
                                             sdate_dim = "syear",
                                             ncores = ncores)
        } else {
          fcst_calibrated <- NULL
        }
      } 
    } else if (recipe$Analysis$Variables$freq %in% c("daily", "daily_mean")) {
      # Daily data calibration using Quantile Mapping
      if (!(method %in% c("qmap"))) {
        error(recipe$Run$logger, retrieve = retrieve,
              paste("Calibration method in the recipe is not available for",
                    "daily data. Only quantile mapping 'qmap is implemented."))
        stop()
      }
      # Calibrate the hindcast
      dim_order <- names(dim(data$hcst$data))
      hcst_calibrated <- CST_QuantileMapping(data$hcst, data$obs,
                                             exp_cor = NULL,
                                             sdate_dim = "syear",
                                             memb_dim = "ensemble",
                                             # window_dim = "time",
                                             method = "QUANT",
                                             ncores = ncores,
                                             na.rm = na.rm,
                                             wet.day = F)
      # Restore dimension order
      browser()
      hcst_calibrated$data <- Reorder(hcst_calibrated$data, dim_order)
      # In the case where anomalies have been computed, calibrate full values
      if (!is.null(data$hcst.full_val)) {
          hcst_full_calibrated <- CST_QuantileMapping(data$hcst.full_val,
                                                      data$obs.full_val,
                                                      exp_cor = NULL,
                                                      sdate_dim = "syear",
                                                      memb_dim = "ensemble",
                                                      method = "QUANT",
                                                      ncores = ncores,
                                                      na.rm = na.rm,
                                                      wet.day = F)
      } else {
        hcst_full_calibrated <- NULL
      }

      if (!is.null(data$fcst)) {
        # Calibrate the forecast
        fcst_calibrated <- CST_QuantileMapping(data$hcst, data$obs,
                                               exp_cor = data$fcst,
                                               sdate_dim = "syear",
                                               memb_dim = "ensemble",
                                               # window_dim = "time",
                                               method = "QUANT",
                                               ncores = ncores,
                                               na.rm = na.rm,
                                               wet.day = F)
        # Restore dimension order
        fcst_calibrated$data <- Reorder(fcst_calibrated$data, dim_order)
      } else {
        fcst_calibrated <- NULL
      }
    }
  }
  info(recipe$Run$logger, retrieve = retrieve, CALIB_MSG)
  .log_memory_usage(recipe$Run$logger, when = "After calibration")
  # Saving
  if (retrieve) {
    if (recipe$Analysis$Workflow$Calibration$save != 'none') {
      info(recipe$Run$logger, retrieve = retrieve,
           "##### START SAVING CALIBRATED DATA #####")
    
      ## TODO: What do we do with the full values?
      recipe$Run$output_dir <- paste0(recipe$Run$output_dir,
                                      "/outputs/Calibration/")
      if ((recipe$Analysis$Workflow$Calibration$save %in% 
           c('all', 'exp_only', 'fcst_only')) && (!is.null(data$fcst))) {
        save_forecast(recipe = recipe, data_cube = fcst_calibrated, type = 'fcst')
      }
      if (recipe$Analysis$Workflow$Calibration$save %in%
          c('all', 'exp_only')) {
        save_forecast(recipe = recipe, data_cube = hcst_calibrated, type = 'hcst')
      }
      if (recipe$Analysis$Workflow$Calibration$save == 'all') {
        save_observations(recipe = recipe, data_cube = data$obs)
      }
    }
  }

  ## TODO: Sort out returns
  return_list <- list(hcst = hcst_calibrated,
                      obs = data$obs,
                      fcst = fcst_calibrated)
  if (!is.null(hcst_full_calibrated)) {
    return_list <- append(return_list,
                          list(hcst.full_val = hcst_full_calibrated,
                               obs.full_val = data$obs.full_val))
  }
  return(return_list)
}
