### Downscaling Module
source('modules/Downscaling/tmp/Interpolation.R')
source('modules/Downscaling/tmp/Intbc.R')
source('modules/Downscaling/tmp/Intlr.R')
source('modules/Downscaling/tmp/Analogs.R')
source('modules/Downscaling/tmp/LogisticReg.R')
source('modules/Downscaling/tmp/Utils.R')

Downscaling <- function(recipe, data, retrieve = TRUE) {
  # Function that downscale the hindcast using the method stated in the
  # recipe. For the moment, forecast must be null.
  #
  # data: list of s2dv_cube objects containing the hcst, obs and fcst.
  # recipe: object obtained when passing the .yml recipe file to read_yaml()
  type <- tolower(recipe$Analysis$Workflow$Downscaling$type)
  
  if (type == "none") {
    hcst_downscal <- data$hcst
    DOWNSCAL_MSG <- "##### NO DOWNSCALING PERFORMED #####"

  } else {

    if (!is.null(data$fcst)) {
      warn(recipe$Run$logger, retrieve = retrieve,
           "The downscaling will be only performed to the hindcast data")
      data$fcst <- NULL
    }
    # Downscaling function params
    int_method     <- tolower(recipe$Analysis$Workflow$Downscaling$int_method)
    bc_method      <- tolower(recipe$Analysis$Workflow$Downscaling$bc_method)
    lr_method      <- tolower(recipe$Analysis$Workflow$Downscaling$lr_method)
    log_reg_method <- tolower(recipe$Analysis$Workflow$Downscaling$log_reg_method)
    target_grid    <- tolower(recipe$Analysis$Workflow$Downscaling$target_grid)
    nanalogs       <- as.numeric(recipe$Analysis$Workflow$Downscaling$nanalogs)
    size           <- recipe$Analysis$Workflow$Downscaling$size   
    
    if (is.null(recipe$Analysis$ncores)) {
      ncores <- 1
    } else {
      ncores <- recipe$Analysis$ncores
    }
 
    #TO DO: add the parametre loocv where it corresponds
    if (is.null(recipe$Analysis$loocv)) {
      loocv <- TRUE
    } else {
      loocv <- recipe$Analysis$loocv
    }

    DOWNSCAL_TYPES <- c("none", "int", "intbc", "intlr", "analogs", "logreg")
    BC_METHODS <- c("quantile_mapping", "bias", "evmos", "mse_min", "crps_min", "rpc-based", "qm")
    LR_METHODS <- c("basic", "large-scale", "4nn")
    LOG_REG_METHODS <- c("ens_mean", "ens_mean_sd", "sorted_members")

    if (!(type %in% DOWNSCAL_TYPES)) {
      stop("Downscaling type in the recipe is not available. Accepted types ",
           "are 'none', 'int', 'intbc', 'intlr', 'analogs', 'logreg'.")
    } 

    if (type == "int") {
      if (is.null(int_method)) {
        stop("Please provide one interpolation method in the recipe.")
      }

      if (is.null(target_grid)) {
        stop("Please provide the target grid in the recipe.")
      }

      # Ensure that observations are in the same grid as experiments
      # Only needed for this method because the others already return the 
      # observations
      latmin <- data$hcst$coords$latitude[1]
      lonmin <- data$hcst$coords$longitude[1]
      latmax <- data$hcst$coords$latitude[length(data$hcst$coords$latitude)]
      lonmax <- data$hcst$coords$longitude[length(data$hcst$coords$longitude)]
      hcst_downscal <- CST_Interpolation(data$hcst, 
                                         points = NULL, 
                                         method_remap = int_method, 
                                         target_grid = target_grid,
                                         lat_dim = "latitude", 
                                         lon_dim = "longitude", 
                                         region = c(lonmin, lonmax, latmin, latmax),
                                         method_point_interp = NULL)

      obs_downscal <- CST_Interpolation(data$obs,
                                        points = NULL,
                                        method_remap = int_method,
                                        target_grid = target_grid,
                                        lat_dim = "latitude",
                                        lon_dim = "longitude",
                                        region = c(lonmin, lonmax, latmin, latmax),
                                        method_point_interp = NULL)

      hcst_downscal$obs <- obs_downscal$exp

      DOWNSCAL_MSG <- "##### DOWNSCALING COMPLETE #####"
    } else if (type == "intbc") {
      if (length(int_method) == 0) {
        stop("Please provide one (and only one) interpolation method in the recipe.")
      }

      if (is.null(bc_method)) {
        stop("Please provide one bias-correction method in the recipe. Accepted ",  
              "methods are 'quantile_mapping', 'bias', 'evmos', 'mse_min', 'crps_min' ",
              "'rpc-based', 'qm'. ")
      }

      if (is.null(target_grid)) {
        stop("Please provide the target grid in the recipe.")
      }

      if (!(bc_method %in% BC_METHODS)) {
        stop(paste0(bc_method, " method in the recipe is not available. Accepted methods ",
             "are 'quantile_mapping', 'bias', 'evmos', 'mse_min', 'crps_min' ",
             "'rpc-based', 'qm'."))
      }
      
        hcst_downscal <- CST_Intbc(data$hcst, data$obs, 
                                   target_grid = target_grid, 
                                   bc_method = bc_method, 
                                   int_method = int_method, 
                                   points = NULL,
                                   method_point_interp = NULL, 
                                   lat_dim = "latitude", 
                                   lon_dim = "longitude",
                                   sdate_dim = "syear", 
                                   member_dim = "ensemble", 
                                   region = NULL, 
                                   ncores = ncores)
	   DOWNSCAL_MSG <- "##### DOWNSCALING COMPLETE #####"
    } else if (type == "intlr") {
      if (length(int_method) == 0) {
        stop("Please provide one (and only one) interpolation method in the recipe.")
      }

      if (is.null(lr_method)) {
        stop("Please provide one linear regression method in the recipe. Accepted ", 
              "methods are 'basic', 'large-scale', '4nn'.")
      }

      if (is.null(target_grid)) {
        stop("Please provide the target grid in the recipe.")
      }

      if (!(lr_method %in% LR_METHODS)) {
        stop(paste0(lr_method, " method in the recipe is not available. Accepted methods ",
             "are 'basic', 'large-scale', '4nn'."))
      }

      # TO DO: add the possibility to have the element 'pred' in 'data'
      if (lr_method == "large-scale") {
        if (is.null(data$pred$data)) {
          stop("Please provide the large scale predictors in the element 'data$pred$data'.")
        }
      } else {
        data$pred$data <- NULL
      }

      hcst_downscal <- CST_Intlr(data$hcst, data$obs, 
                                 lr_method = lr_method, 
                                 target_grid = target_grid, 
                                 points = NULL, 
                                 int_method = int_method,
                                 method_point_interp = NULL, 
                                 predictors = data$pred$data, 
                                 lat_dim = "latitude", 
                                 lon_dim = "longitude",
                                 sdate_dim = "syear", 
                                 time_dim = "time", 
                                 member_dim = "ensemble",
                                 large_scale_predictor_dimname = 'vars',
                                 loocv = loocv, 
                                 region = NULL, 
                                 ncores = ncores)

      DOWNSCAL_MSG <- "##### DOWNSCALING COMPLETE #####"
    } else if (type == "analogs") {
      
      if (is.null(nanalogs)) {
        warning("The number of analogs for searching has not been provided in the ",
                "recipe. Setting it to 3.")
        nanalogs <- 3
      }

      if (!is.null(size) & recipe$Analysis$Variables$freq == "monthly_mean") {
        size <- NULL
        warning("Size is set to NULL. ",
                "It must be NULL for the monthly input data.")
      }

      if (!is.null(size)) {
        dum <- data$obs$data ## keep obs data before the window process to provide it later as the  output
        data$obs$data <- .generate_window(data$obs$data,
                                          sdate_dim = 'syear',
                                          time_dim = 'time',
                                          loocv = TRUE,
                                          size = size)
        data$obs$data <- Apply(data$obs$data,
                               target_dims="window",
                               fun=function (x) x[!is.na(x)])$output1
      }

      hcst_downscal <- CST_Analogs(data$hcst, data$obs, 
                                   grid_exp = data$hcst$attrs$source_files[
                                              which(!is.na(data$hcst$attrs$source_files))[1]], 
                                   nanalogs = nanalogs, 
                                   fun_analog = "wmean", 
                                   lat_dim = "latitude", 
                                   lon_dim = "longitude",
                                   sdate_dim = "syear", 
                                   time_dim = "time", 
                                   member_dim = "ensemble",
                                   region = NULL, 
                                   return_indices = FALSE,
                                   loocv_window = loocv, 
                                   ncores = ncores)
      
      if (!is.null(size)) {
      hcst_downscal$obs$data <- Apply(dum, target_dims=c("time", "smonth"),
                                      function (x) {x[1:(dim(data$hcst$data)["time"]), 2]},
                                      ncores = ncores,
                                      output_dims = "time")$output1 ## 2nd month is the target month 
      }

      DOWNSCAL_MSG <- "##### DOWNSCALING COMPLETE #####"
    } else if (type == "logreg") {
     
      if (length(int_method) == 0) {
        stop("Please provide one (and only one) interpolation method in the recipe.")
      }
 
      if (is.null(log_reg_method)) {
        stop("Please provide one logistic regression method in the recipe. Accepted ", 
              "methods are 'ens_mean', 'ens_mean_sd', 'sorted_members'.")
      }

      if (is.null(target_grid)) {
        stop("Please provide the target grid in the recipe.")
      }

      # Since we are forcing to create three categories, and applying cross-validation,
      # we need at least six years of data for the logistic regression function to not 
      # crash
      if (dim(data$hcst$data)[names(dim(data$hcst$data)) == "syear"] <= 5) {
        stop("The number of start dates is insufficient for the logisitic regression method. ",
             "Please provide six or more.") 
      }
 
      if (!(log_reg_method %in% LOG_REG_METHODS)) {
        stop(paste0(log_reg_method, " method in the recipe is not available. Accepted methods ",
             "are 'ens_mean', 'ens_mean_sd', 'sorted_members'."))
      }

      hcst_downscal <- CST_LogisticReg(data$hcst, data$obs, 
                                       target_grid = target_grid, 
                                       int_method = int_method, 
                                       log_reg_method = log_reg_method,
                                       probs_cat = c(1/3,2/3), 
                                       return_most_likely_cat = FALSE, 
                                       points = NULL,
                                       method_point_interp = NULL, 
                                       lat_dim = "latitude", 
                                       lon_dim = "longitude", 
                                       sdate_dim = "syear",
                                       member_dim = "ensemble", 
                                       region = NULL, 
                                       loocv = loocv, 
                                       ncores = ncores)

      DOWNSCAL_MSG <- "##### DOWNSCALING COMPLETE #####"
    }
  } 
  info(recipe$Run$logger, retrieve = retrieve, DOWNSCAL_MSG)

  # Saving
  if (retrieve) {
    if (recipe$Analysis$Workflow$Downscaling$save != 'none') {
      info(recipe$Run$logger, retrieve = retrieve,
           "##### START SAVING DOWNSCALED DATA #####")
    }
    recipe$Run$output_dir <- paste0(recipe$Run$output_dir, 
                                    "/outputs/Downscaling/")
    # if ((recipe$Analysis$Workflow$Downscaling$save %in% 
    #      c('all', 'exp_only', 'fcst_only')) && (!is.null(data$fcst))) {
    #   save_forecast(recipe = recipe, data_cube = data$fcst, type = 'fcst')
    # }
    if (recipe$Analysis$Workflow$Downscaling$save %in% c('all', 'exp_only')) {
      save_forecast(recipe = recipe,
                    data_cube = hcst_downscal$exp,
                    type = 'hcst')
    }
    if (recipe$Analysis$Workflow$Downscaling$save == 'all') {
      save_observations(recipe = recipe,
                        data_cube = hcst_downscal$obs)
    }
  }

  return(list(hcst = hcst_downscal$exp, obs = hcst_downscal$obs, fcst = NULL))
}
