### 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)
  fcst_only <- recipe$Analysis$Workflow$Downscaling$fcst_only

  if (is.null(data$fcst) && is.null(data$fcstL)) {
    # In this case, fcst_only is set to TRUE because the forecast call will
    # return the hindcast, so that the hindcast is not computed twice.
    fcst_only <- TRUE
  } else if (!isTRUE(fcst_only)) {
    fcst_only <- FALSE
  }

  if (type == "none") {
    DOWNSCAL_MSG <- "##### METHOD SET TO 'none': NO DOWNSCALING PERFORMED #####"
  } else {
    # 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    <- recipe$Analysis$Workflow$Downscaling$target_grid
    nanalogs       <- as.numeric(recipe$Analysis$Workflow$Downscaling$nanalogs)
    size           <- recipe$Analysis$Workflow$Downscaling$size   
    probs_cat      <- sapply(unlist(recipe$Analysis$Workflow$Downscaling$probs_cat),
                             function(frac) eval(parse(text = frac)))
    if (tolower(recipe$Analysis$Horizon) == "seasonal") {
      archive <- read_yaml("conf/archive_seasonal.yml")
    } else if (tolower(recipe$Analysis$Horizon) == "subseasonal") {
      archive <- read_yaml("conf/archive_subseasonal.yml")
    } else if (tolower(recipe$Analysis$Horizon) == "decadal") {
      archive <- read_yaml("conf/archive_decadal.yml")
    }
    system_info <- archive[[recipe$Run$filesystem]]$System[[recipe$Analysis$Datasets$System$name]]
    
    if(!is.numeric(probs_cat)) {
      probs_cat <- c(1/3, 2/3) 
    }

    if (is.null(recipe$Analysis$ncores)) {
      ncores <- 1
    } else {
      ncores <- recipe$Analysis$ncores
    }
 
    loocv_hcst <- TRUE # loocv for the hcst dwn, in case the fcst dwn is also requested.
    if (!is.null(data$fcst) | (!is.null(data$fcstL) & type == "analogs")) {
      loocv <- FALSE
    } else  {
      loocv <- TRUE	
    }

    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", "9nn")
    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'."))
      }

      downscal <- CST_Intbc(data$hcst, data$obs, data$fcst, 
                            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)

      if (!fcst_only & !is.null(data$fcst)) {
        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', '9nn'.")
      }

      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', '9nn'."))
      }

      # 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
      }

      downscal <- CST_Intlr(data$hcst, data$obs, data$fcst,
                            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)

      if (!fcst_only & !is.null(data$fcst)) {
        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_hcst, 
                                   region = NULL, 
                                   ncores = ncores)
      }
      DOWNSCAL_MSG <- "##### DOWNSCALING COMPLETE #####"

    } else if (type == "analogs") {

      if ((!is.null(data$hcstL) | !is.null(data$fcstL)) & is.null(data$obsL)) {
        stop("Large scale model field (i.e., data$hcstL or data$fcstL) is provided. ",
             "Observation must also be provided for the large scale variable.")
      }
      
      if (is.null(data$hcstL) & is.null(data$fcstL) & !is.null(data$obsL)) {
        stop("Large scale observational field (i.e., data$obsL) is provided. ",
             "Model data must also be provided for the large scale variable.")
      }

      if (!is.null(data$hcstL) & !is.null(data$obsL) & !is.null(data$fcst) & is.null(data$fcstL)) {
        stop("Large scale observational field (i.e., data$obsL) is provided. Please also ",
             "provide large scale forecast data (i.e., data$fcstL) ",
             "or set data$fcst to NULL.")
      }

      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" | 
                            recipe$Analysis$Variables$freq == "weekly_mean")) {
        size <- NULL
        warning("Size is set to NULL. ",
                "It must be NULL for the weekly or monthly input data.")
      }

      if (!is.null(size)) {
        ## keep obs data before the window process to provide it later as the  output
        dum <- data$obs$data
        data$obs$data <- .generate_window(data$obs$data,
                                          sdate_dim = 'syear',
                                          time_dim = 'time',
                                          loocv = TRUE,
                                          size = size,
					  ncores = ncores)

        if (!is.null(data$hcstL)) {
          data$obsL$data <- .generate_window(data$obsL$data,
                                             sdate_dim = 'syear',
                                             time_dim = 'time',
                                             loocv = TRUE,
                                             size = size,
                                             ncores = ncores)
        }
      }

      ## Priority: fcstL, fcst, hcstL, hcst
      exp <- data$hcst
      if (!is.null(data$hcstL)) {
        exp <- data$hcstL
        if (is.null(data$fcst)) {
          print("Provided large scale variable is used as the predictor.")
        }
      }
      if (!is.null(data$fcst)) {
        exp <- data$fcst
      }  
      if (!is.null(data$fcstL)) {
        exp <- data$fcstL
        if (!is.null(data$hcstL) & !is.null(data$fcst)) {
          print("Provided large scale variable is used as the predictor.")
        }
      }
      if (!is.null(system_info)) {
        grid_exp <- system_info$reference_grid
      } else {
        grid_exp <- exp$attrs$source_files[which(!is.na(exp$attrs$source_files))[1]]
      }

      downscal <- CST_Analogs(exp, data$obs, 
                              grid_exp = grid_exp,
                              obsL = data$obsL,
                              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,
	                      metric = "dist", 
                              ncores = ncores)

      if (!fcst_only & (!is.null(data$fcst) | !is.null(data$fcstL))) {
        ## Priority: hcstL, hcst
        exp <- data$hcst
        if (!is.null(data$hcstL)) {
          exp <- data$hcstL
        }
        hcst_downscal <- CST_Analogs(exp, data$obs,
                                     grid_exp = grid_exp,
                                     obsL = data$obsL,
                                     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_hcst,
                                     metric = "dist",
                                     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'."))
      }

      downscal <- CST_LogisticReg(data$hcst, data$obs, data$fcst, 
                                  target_grid = target_grid, 
                                  int_method = int_method, 
                                  log_reg_method = log_reg_method,
                                  probs_cat = probs_cat, 
                                  return_most_likely_cat = FALSE, 
                                  points = NULL,
                                  method_point_interp = NULL, 
                                  lat_dim = "latitude", 
                                  lon_dim = "longitude", 
                                  sdate_dim = "syear",
                                  time_dim = "time",
                                  member_dim = "ensemble", 
                                  region = NULL, 
                                  loocv = loocv, 
                                  ncores = ncores)

      if (!fcst_only && !is.null(data$fcst)) {
        hcst_downscal <- CST_LogisticReg(data$hcst, data$obs,
                                         target_grid = target_grid,
                                         int_method = int_method,
                                         log_reg_method = log_reg_method,
                                         probs_cat = probs_cat,
                                         return_most_likely_cat = FALSE,
                                         points = NULL,
                                         method_point_interp = NULL,
                                         lat_dim = "latitude",
                                         lon_dim = "longitude",
                                         sdate_dim = "syear",
                                         time_dim = "time",
                                         member_dim = "ensemble",
                                         region = NULL,
                                         loocv = loocv_hcst,
                                         ncores = ncores)
      }
      DOWNSCAL_MSG <- "##### DOWNSCALING COMPLETE #####"

    }
  } 
  print(DOWNSCAL_MSG)

  # Assign the results to the corresponding object
  data$obs <- downscal$obs
  if (is.null(data$fcst) && is.null(data$fcstL)) {
    # If there is no fcst, downscal$exp is the downscaled hcst
    data$hcst <- downscal$exp
    data$fcst <- NULL
  } else {
    # If there is a fcst, downscal$exp is the downscaled fcst
    data$fcst <- downscal$exp
    if (!fcst_only) {
      data$hcst <- hcst_downscal$exp
    } else {
      data$hcst <- NULL
      warn(recipe$Run$logger,
           paste("Only $obs and $fcst have been downscaled; $hcst will",
                 "be NULL and some modules might return errors."))
    }
  }

  # Saving
  if (retrieve) {
    if (recipe$Analysis$Workflow$Downscaling$save != 'none') {
      info(recipe$Run$logger, "##### 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')) && (!is.null(data$hcst))) {
      save_forecast(recipe = recipe, data_cube = data$hcst, type = 'hcst')
    }
    if (recipe$Analysis$Workflow$Downscaling$save == 'all') {
      save_observations(recipe = recipe, data_cube = data$obs)
    }
  }

  return(list(hcst = data$hcst, obs = data$obs, fcst = data$fcst))
}
