# This module should calculate verification metrics at any time of the workflow
# It should implement different verification metrics:
# - FRPSS and RPSS
# - FCRPSS and CRPSS
# - enscorr
# - bias
# - reliability diagram
# - ask Carlos which decadal metrics he is currently using

source("modules/Skill/R/compute_quants.R")
source("modules/Skill/R/compute_probs.R")
source("modules/Skill/R/s2s.metrics.R")
source("modules/Saving/R/drop_dims.R")
## TODO: Remove when they are included in s2dv
## TODO: Add pending functions
source("modules/Skill/R/RPS_clim.R")
source("modules/Skill/R/CRPS_clim.R")
source("modules/Skill/R/tmp/GetProbs.R")
source("modules/Skill/R/tmp/RPS.R")
source("modules/Skill/R/tmp/CRPS.R")
source("modules/Skill/R/tmp/SprErr.R") # https://earth.bsc.es/gitlab/es/s2dv/-/issues/115
source("modules/Skill/R/tmp/Bias.R") # https://earth.bsc.es/gitlab/es/s2dv/-/issues/118
source("modules/Skill/R/tmp/Eno.R")

Skill <- function(recipe, data, agg = 'global', retrieve = TRUE,
                  nchunks = nchunks) {
  
  # data$hcst: s2dv_cube containing the hindcast
  # obs: s2dv_cube containing the observations
  # recipe: auto-s2s recipe as provided by read_yaml
  
  ## TODO: Adapt time_dims to subseasonal case
  ## TODO: Add dat_dim
  ## TODO: Refine error messages
  ## TODO: Add check to see if anomalies are provided (info inside s2dv_cube)

#   if (recipe$Analysis$Workflow$Anomalies$compute) {
#     if (is.null(clim_data$hcst) || is.null(clim_obs)) {
#       warn(recipe$Run$logger, "Anomalies have been requested in the recipe,
#          but the climatologies have not been provided in the 
#          compute_skill_metrics call. Be aware that some metrics like the 
#          Mean Bias may not be correct.")
#     }
#   } else {
#     warn(recipe$Run$logger, "Anomaly computation was not requested in the
#        recipe. Be aware that some metrics, such as the CRPSS may not be 
#        correct.")
#   }
  time_dim <- 'syear'
  memb_dim <- 'ensemble'
  metrics <- tolower(recipe$Analysis$Workflow$Skill$metric)
  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
  }
  ## TODO: Move this to the recipe checker
  if (is.null(recipe$Analysis$Workflow$Skill$cross_validation)) {
  #  warn(recipe$Run$logger,
  #       "cross_validation parameter not defined, setting it to FALSE.")
    cross.val <- FALSE
  } else {
    cross.val <- recipe$Analysis$Workflow$Skill$cross_validation
  }
  skill_metrics <- list()
  for (metric in strsplit(metrics, ", | |,")[[1]]) { 
    # Whether the fair version of the metric is to be computed
    if (metric %in% c('frps', 'frpss', 'bss10', 'bss90',
                      'fcrps', 'fcrpss')) {
      Fair <- T
    } else {
      Fair <- F
    }
    # Whether to compute correlation for the ensemble mean or for each member
    if (metric == 'corr_individual_members') {
      memb <- T
    } else if (metric == 'enscorr') {
      memb <- F
    } 
    # Ranked Probability Score and Fair version
    if (metric %in% c('rps', 'frps')) {
      skill <- RPS(data$hcst$data, data$obs$data,
                   time_dim = time_dim,
                   memb_dim = memb_dim,
                   Fair = Fair,
                   cross.val = cross.val,
                   return_mean = TRUE,
                   ncores = ncores)
      skill <- .drop_dims(skill)
      skill_metrics[[ metric ]] <- skill
    # RPS_clim
    } else if (metric == 'rps_clim') {
      skill <- Apply(list(data$obs$data),
                     target_dims = c(time_dim, memb_dim),
                     cross.val = cross.val,
                     return_mean = TRUE,
                     fun = RPS_clim)$output1
      skill <- .drop_dims(skill)
      skill_metrics[[metric]] <- skill
    # RPS_syear and FRPS_syear
    } else if (metric %in% c('rps_syear', 'frps_syear')) {
      skill <- RPS(data$hcst$data, data$obs$data,
                   time_dim = time_dim,
                   memb_dim = memb_dim,
                   Fair = Fair,
                   cross.val = cross.val,
                   return_mean = FALSE,
                   ncores = ncores)
      skill <- .drop_dims(skill)
      skill_metrics[[ metric ]] <- skill ## temp
    # RPS_clim_syear
    } else if (metric == 'rps_clim_syear') { 
      skill <- Apply(list(data$obs$data),
                        target_dims = c(time_dim, memb_dim),
                        cross.val = cross.val,
                        fun = RPS_clim,
                        return_mean = FALSE, output_dims = 'syear')$output1
      skill <- .drop_dims(skill)
      skill_metrics[[ metric ]] <- skill ## temp
    # Ranked Probability Skill Score and Fair version
    } else if (metric %in% c('rpss', 'frpss')) {
      # skill <- RPSS(data$hcst$data, data$obs$data,
      skill <- RPSS(data$hcst$data, data$obs$data,
		    time_dim = time_dim,
		    memb_dim = memb_dim,
          	    Fair = Fair,
		    ncores = ncores)
      skill <- lapply(skill, function(x) {
                      .drop_dims(x)})
      skill_metrics[[ metric ]] <- skill$rpss
      skill_metrics[[ paste0(metric, "_significance") ]] <- skill$sign
    # Brier Skill Score - 10th percentile
    } else if (metric == 'bss10') {
      skill <- RPSS(data$hcst$data, data$obs$data,
                    time_dim = time_dim, 
                    memb_dim = memb_dim,
                    prob_thresholds = 0.1,
                    cross.val = cross.val,
                    Fair = Fair,
                    ncores = ncores)
      skill <- lapply(skill, function(x) {
                      .drop_dims(x)})
      skill_metrics[[ metric ]] <- skill$rpss
      skill_metrics[[ paste0(metric, "_significance") ]] <- skill$sign
    # Brier Skill Score - 90th percentile
    } else if (metric == 'bss90') {
      skill <- RPSS(data$hcst$data, data$obs$data,
                    time_dim = time_dim,
                    memb_dim = memb_dim,
                    prob_thresholds = 0.9,
                    cross.val = cross.val,
                    Fair = Fair,
                    ncores = ncores)
      skill <- lapply(skill, function(x) {
                      .drop_dims(x)})
      skill_metrics[[ metric ]] <- skill$rpss
      skill_metrics[[ paste0(metric, "_significance") ]] <- skill$sign
    # CRPS and FCRPS
    } else if (metric %in% c('crps', 'fcrps')) {
      skill <- CRPS(data$hcst$data, data$obs$data,
                    time_dim = time_dim,
                    memb_dim = memb_dim,
                    Fair = Fair,
                    return_mean = TRUE, 
                    ncores = ncores)
      skill <- .drop_dims(skill)
      skill_metrics[[ metric ]] <- skill
    # CRPS_clim
    } else if (metric %in% c('crps_clim')) {    
      skill <- Apply(list(data$obs$data), target_dims = time_dim,
                     fun = CRPS_clim, memb_dim = memb_dim,
                     clim.cross.val = cross.val, 
                     return_mean = TRUE)$output1
      skill <- .drop_dims(skill)
      skill_metrics[[ metric ]] <- skill
    # CRPS_syear and FCRPS_syear
    } else if (metric %in% c('crps_syear', 'fcrps_syear')) {
      skill <- CRPS(data$hcst$data, data$obs$data,
                    time_dim = time_dim,
                    memb_dim = memb_dim,
                    Fair = Fair,
                    return_mean = FALSE, 
                    ncores = ncores)
      skill <- .drop_dims(skill)
      skill_metrics[[ metric ]] <- skill
    # CRPS_clim_syear
    } else if (metric %in% c('crps_clim_syear')) {
      skill <- Apply(list(data$obs$data), target_dims = time_dim,
                         fun = CRPS_clim, memb_dim = memb_dim,
                         clim.cross.val = cross.val,
                         return_mean = FALSE)$output1
      skill <- .drop_dims(skill)
      skill_metrics[[metric]] <- skill
    # CRPSS and FCRPSS
    } else if (metric %in% c('crpss', 'fcrpss')) {
      skill <- CRPSS(data$hcst$data, data$obs$data,
                     time_dim = time_dim,
                     memb_dim = memb_dim,
                     clim.cross.val = cross.val,
                     Fair = Fair,
                     ncores = ncores)
      skill <- lapply(skill, function(x) {
                      .drop_dims(x)})
      skill_metrics[[ metric ]] <- skill$crpss
      skill_metrics[[ paste0(metric, "_significance") ]] <- skill$sign
    } else if (metric == 'rms') {
      hcst_mean <- Apply(list(data$hcst$data), target_dims = memb_dim,
                         fun = mean, na.rm = na.rm, ncores = ncores)$output1
      hcst_mean <- InsertDim(hcst_mean, pos = 1, len = 1, name = memb_dim)
      skill <- RMS(exp = hcst_mean, 
                   obs = data$obs$data, 
                   time_dim = time_dim, dat_dim = NULL, comp_dim = NULL,
                   limits = NULL, conf = FALSE, alpha = 0.05, ncores = ncores)
      skill <- lapply(skill, function(x) {
                      .drop_dims(x)})
      skill_metrics[[ metric ]] <- skill$rms
    # Mean bias (climatology) 
    } else if (metric == 'mean_bias') {
      ## TODO: Eliminate option to compute from anomalies
      # Compute from full field
      if ((!is.null(data$hcst.full_val)) && (!is.null(data$obs.full_val)) && 
          (recipe$Analysis$Workflow$Anomalies$compute)) {
        skill <- Bias(data$hcst.full_val$data, data$obs.full_val$data,
                      time_dim = time_dim,
                      memb_dim = memb_dim,
                      alpha = 0.05,
                      ncores = ncores)
      } else {
        skill <- Bias(data$hcst$data, data$obs$data,
                      time_dim = time_dim,
                      memb_dim = memb_dim,
                      alpha = 0.05,
                      ncores = ncores)
      }
      skill <- lapply(skill, function(x) {
        .drop_dims(x)})
      skill_metrics[[ metric ]] <- skill$bias
      skill_metrics[[ paste0(metric, "_significance") ]] <- skill$sign
    # Mean bias skill score
    } else if (metric == 'mean_bias_ss') {
      if ((!is.null(data$hcst.full_val)) && (!is.null(data$obs.full_val)) && 
          (recipe$Analysis$Workflow$Anomalies$compute)) {
        skill <- AbsBiasSS(data$hcst.full_val$data, data$obs.full_val$data,
                           time_dim = time_dim,
                           memb_dim = memb_dim,
                           ncores = ncores)
      } else {
        skill <- AbsBiasSS(data$hcst$data, data$obs$data,
                           time_dim = time_dim,
                           memb_dim = memb_dim,
                           ncores = ncores)
      }
      skill <- lapply(skill, function(x) {
                      .drop_dims(x)})
      skill_metrics[[ metric ]] <- skill$biasSS
      skill_metrics[[ paste0(metric, "_significance") ]] <- skill$sign
    # Ensemble mean correlation
    } else if (metric %in% c('enscorr', 'corr_individual_members')) {
      ## TODO: Return significance
      ## TODO: Implement option for Kendall and Spearman methods?
      skill <- s2dv::Corr(data$hcst$data, data$obs$data, 
                          dat_dim = 'dat', 
                          time_dim = time_dim, 
                          method = 'pearson', 
                          memb_dim = memb_dim,
                          memb = memb,
                          conf = F,
                          pval = F,
                          sign = T,
                          alpha = 0.05,
                          ncores = ncores)
      skill <- lapply(skill, function(x) {
                      .drop_dims(x)})
      skill_metrics[[ metric ]] <- skill$corr
      skill_metrics[[ paste0(metric, "_significance") ]] <- skill$sign
    } else if (metric == 'rmsss') {
      # Compute RMSSS
      skill <- RMSSS(data$hcst$data, data$obs$data,
                     dat_dim = 'dat',
                     time_dim = time_dim,
                     memb_dim = memb_dim,
                     pval = FALSE,
                     sign = TRUE,
                     sig_method = 'Random Walk',
                     ncores = ncores)
      # Compute ensemble mean and modify dimensions
      skill <- lapply(skill, function(x) {
                      .drop_dims(x)})
      skill_metrics[[ metric ]] <- skill$rmsss
      skill_metrics[[ paste0(metric, "_significance") ]] <- skill$sign
    } else if (metric == 'mse') {
      skill <- MSE(data$hcst$data, data$obs$data,
                   time_dim = time_dim,
                   memb_dim = memb_dim,
                   dat_dim = 'dat',
                   comp_dim = NULL,
                   conf = FALSE,
                   ncores = ncores)
      skill <- lapply(skill, function(x) {
                      .drop_dims(x)})
      skill_metrics[[ metric ]] <- skill$mse
      skill_metrics[[ paste0(metric, "_significance") ]] <- skill$sign
    } else if (metric == 'msss') {
      skill <- MSSS(data$hcst$data, data$obs$data,
                    time_dim = time_dim,
                    memb_dim = memb_dim,
                    dat_dim = 'dat',
                    sign = TRUE,
                    ncores = ncores)
      skill <- lapply(skill, function(x) {
                      .drop_dims(x)})
      skill_metrics[[ metric ]] <- skill$msss
      skill_metrics[[ paste0(metric, "_significance") ]] <- skill$sign
    } else if (metric == 'enssprerr') {
      .Eno <- s2dv:::.Eno
      skill <- SprErr(data$hcst$data, data$obs$data,
                      time_dim = time_dim,
                      memb_dim = memb_dim,
                      dat_dim = 'dat',
                      pval = FALSE, 
                      sign = TRUE, 
                      alpha = 0.05,
                      na.rm = na.rm, 
                      ncores = ncores)
      skill <- lapply(skill, function(x) {
        .drop_dims(x)})
      skill_metrics[[ metric ]] <- skill$ratio
      skill_metrics[[ paste0(metric, "_significance") ]] <- skill$sign
    # SpecsVerification metrics  
    } else if (grepl("specs", metric, fixed = TRUE)) {
      # Compute SpecsVerification version of the metrics
      ## Retain _specs in metric name for clarity
      metric_name <- (strsplit(metric, "_"))[[1]][1] # Get metric name
      ## TODO: Move this to the recipe checker
      # if (!(metric_name %in% c('frpss', 'frps', 'bss10', 'bss90', 'enscorr',
      #   		       'rpss'))) {
      #   warn(recipe$Run$logger,
      #        "Some of the requested SpecsVerification metrics are not available.")
      # }
      capture.output(
      skill <- Compute_verif_metrics(data$hcst$data, data$obs$data,
                                     skill_metrics = metric_name,
                                     verif.dims=c("syear", "sday", "sweek"),
                                     na.rm = na.rm,
                                     ncores = ncores)
      )
      skill <- .drop_dims(skill)
      if (metric_name == "frps") {
        # Compute yearly mean for FRPS
        skill <- colMeans(skill, dims = 1)
      }
      skill_metrics[[ metric ]] <- skill
    }
  }
  ## TODO: Create list of returns with data and metadata
  info(recipe$Run$logger, retrieve = retrieve,
       "##### SKILL METRIC COMPUTATION COMPLETE #####")
  .log_memory_usage(recipe$Run$logger, when = "After skill metric computation")
  # Save outputs
  if (retrieve) {
    skill_metrics$metadata <- data$hcst
    skill_metrics$metadata$data <- NULL
    ## TODO: Handle the output differently to avoid changing the recipe?
    if (recipe$Analysis$Workflow$Skill$save != 'none') {
      info(recipe$Run$logger, retrieve = retrieve,
           "##### START SAVING SKILL METRIC #####")
    }
    # recipe$Run$output_dir <- paste0(recipe$Run$output_dir,
    #                                 "/outputs/Skill/")
    # Separate 'corr' from the rest of the metrics because of extra 'ensemble' dim
    if (recipe$Analysis$Workflow$Skill$save == 'all') {
      corr_metric_names <- grep("^corr_individual_members",
                                names(skill_metrics))
      if (length(corr_metric_names) == 0) {
        save_metrics(recipe = recipe, metrics = skill_metrics,
                     module = "Skill", agg = agg)
      } else {
        # Save corr
        if (length(skill_metrics[corr_metric_names]) > 0) {
          save_corr(recipe = recipe,
                    skill = skill_metrics[corr_metric_names],
                    data_cube = skill_metrics$metadata,
                    agg = agg)
        }
        # Save other skill metrics
        if (length(skill_metrics[-corr_metric_names]) > 0) {
          save_metrics(recipe = recipe,
                       metrics = skill_metrics[-corr_metric_names],
                       module = "Skill",
                       agg = agg)
        }
      }
    }
  } else {
    # Retrieve original dimensions
    dimensions <- names(dim(skill_metrics[[1]]))
    # Bind list into an array
    skill_metrics <- abind(skill_metrics, along = 0)
    # Retrieve dimname attributes
    metric_names <- dimnames(skill_metrics)[[1]]
    tmp_dir <- paste0(recipe$Run$output_dir, "/outputs/tmp/Skill/")
    # Save metric names (only done once)
    if (!dir.exists(tmp_dir)) {
      dir.create(tmp_dir, recursive = TRUE)
      saveRDS(metric_names, paste0(tmp_dir, "metric_names.Rds"))
    }
    # Add 'metric' dimension to array and restore attributes
    names(dim(skill_metrics)) <- c("metric", dimensions)
    # Save s2dv_cube metadata
    metadata_filename <- "skill_metadata"
    for (chunk in names(nchunks)) {
      metadata_filename <- paste0(metadata_filename, "_", chunk, "_", nchunks[[chunk]])
    }
    metadata_filename <- paste0(metadata_filename, ".Rds")
    data$hcst$data <- NULL
    saveRDS(data$hcst, paste0(tmp_dir, metadata_filename))
  }
  # Return results
  return(skill_metrics)
}

Probabilities <- function(recipe, data, retrieve = TRUE) {
  ## TODO: Do hcst and fcst at the same time
  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
  }
  
  named_probs <- list()
  named_probs_fcst <- list()
  named_quantiles <- list()
  
  if (is.null(recipe$Analysis$Workflow$Probabilities$percentiles)) {
    error(recipe$Run$logger, retrieve = retrieve,
          paste("Quantiles and probability bins have been requested, but no",
                "thresholds are provided in the recipe."))
    stop()
  } else {
    for (element in recipe$Analysis$Workflow$Probabilities$percentiles) {
      # Parse thresholds in recipe
      thresholds <- sapply(element, function (x) eval(parse(text = x)))
      quants <- compute_quants(data$hcst$data, thresholds,
                               ncores = ncores,
                               na.rm = na.rm)
      probs <- compute_probs(data$hcst$data, quants,
                             ncores = ncores,
                             na.rm = na.rm)

      for (i in seq(1:dim(quants)['bin'][[1]])) {
        named_quantiles <- append(named_quantiles,
                                  list(ClimProjDiags::Subset(quants,
                                                             'bin', i)))
        names(named_quantiles)[length(named_quantiles)] <- paste0("percentile_",
                                            as.integer(thresholds[i]*100))
      }
      for (i in seq(1:dim(probs)['bin'][[1]])) {
        if (i == 1) {
          name_i <- paste0("prob_b", as.integer(thresholds[1]*100))
        } else if (i == dim(probs)['bin'][[1]]) {
          name_i <- paste0("prob_a", as.integer(thresholds[i-1]*100))
        } else {
          name_i <- paste0("prob_", as.integer(thresholds[i-1]*100), "_to_",
                           as.integer(thresholds[i]*100))
        }
        named_probs <- append(named_probs,
                              list(ClimProjDiags::Subset(probs,
                                                         'bin', i)))
        names(named_probs)[length(named_probs)] <- name_i
      }

      # Compute fcst probability bins
      if (!is.null(data$fcst)) {
        probs_fcst <- compute_probs(data$fcst$data, quants,
                                    ncores = ncores,
                                    na.rm = na.rm)

        for (i in seq(1:dim(probs_fcst)['bin'][[1]])) {
          if (i == 1) {
            name_i <- paste0("prob_b", as.integer(thresholds[1]*100))
          } else if (i == dim(probs_fcst)['bin'][[1]]) {
            name_i <- paste0("prob_a", as.integer(thresholds[i-1]*100))
          } else {
            name_i <- paste0("prob_", as.integer(thresholds[i-1]*100), "_to_",
                           as.integer(thresholds[i]*100))
          }
          named_probs_fcst <- append(named_probs_fcst,
                                   list(ClimProjDiags::Subset(probs_fcst,
                                                              'bin', i)))
          names(named_probs_fcst)[length(named_probs_fcst)] <- name_i
        }
      }
    }
    # Rearrange dimensions and return probabilities
    named_probs <- lapply(named_probs, function(x) {.drop_dims(x)})
    named_quantiles <- lapply(named_quantiles, function(x) {.drop_dims(x)})
    if (!is.null(data$fcst)) {
      fcst_years <- dim(data$fcst$data)[['syear']]
      named_probs_fcst <- lapply(named_probs_fcst, function(x) {.drop_dims(x)})
      #                           function(x) {Subset(x,
      #                                               along = 'syear', 
      #                                               indices = 1:fcst_years,
      #                                               drop = 'non-selected')})
      results <- list(probs = named_probs,
                      probs_fcst = named_probs_fcst,
                      percentiles = named_quantiles)
    } else {
      results <- list(probs = named_probs,
                      percentiles = named_quantiles)
    }
    
    info(recipe$Run$logger, retrieve = retrieve,
	 "##### PERCENTILES AND PROBABILITY CATEGORIES COMPUTED #####")
    .log_memory_usage(recipe$Run$logger, when = "After anomaly computation")
    # Save outputs
    if (retrieve) {
      if (recipe$Analysis$Workflow$Probabilities$save != 'none') {
        info(recipe$Run$logger, retrieve = retrieve,
             "##### START SAVING PERCENTILES AND PROBABILITY CATEGORIES #####")
      
        recipe$Run$output_dir <- paste0(recipe$Run$output_dir,
                                        "/outputs/Skill/")
        # Save percentiles
        if (recipe$Analysis$Workflow$Probabilities$save %in%
            c('all', 'percentiles_only')) {
          save_percentiles(recipe = recipe, percentiles = results$percentiles,
                           data_cube = data$hcst)
        }
        # Save probability bins
        if (recipe$Analysis$Workflow$Probabilities$save %in%
            c('all', 'bins_only')) {
          save_probabilities(recipe = recipe, probs = results$probs,
                             data_cube = data$hcst, type = "hcst")
          if (!is.null(results$probs_fcst)) {
            save_probabilities(recipe = recipe, probs = results$probs_fcst,
                               data_cube = data$fcst, type = "fcst")
          }
        }
      }
    }
    # Return results
    return(results)
  }
}

