# 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
source("modules/Skill/R/RPS_clim.R")
source("modules/Skill/R/CRPS_clim.R")
source("modules/Skill/R/tmp/GetProbs.R")
## TODO: Remove in the next release
source("modules/Skill/compute_skill_metrics.R")
source("modules/Skill/compute_probabilities.R")

source("SUNSET_PROV/R/Prov.Skill.R")
source("SUNSET_PROV/R/SunsetProv.Skill.R")
source("SUNSET_PROV/R/provenance-helpers.R")
source("SUNSET_PROV/R/igraph-helpers.R")

Skill <- function(recipe, data, agg = 'global') {

  # 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.")
#   }
  #Extract provenance graph
  graph.prov <- data$graph.prov

  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
  }
  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,
                   ncores = ncores)
      skill <- .drop_dims(skill)
      skill_metrics[[ metric ]] <- skill
      rps_clim <- Apply(list(data$obs$data),
                             target_dims = c(time_dim, memb_dim),
                             cross.val = cross.val,
                             RPS_clim)$output1
      rps_clim <- .drop_dims(rps_clim)
      skill_metrics[[paste0(metric, "_clim")]] <- rps_clim
    # Ranked Probability Skill Score and Fair version
    } else if (metric %in% c('rpss', 'frpss')) {
      skill <- RPSS(data$hcst$data, data$obs$data,
                    time_dim = time_dim,
                    memb_dim = memb_dim,
                    Fair = Fair,
                    cross.val = cross.val,
                    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,
                    ncores = ncores)
      skill <- .drop_dims(skill)
      skill_metrics[[ metric ]] <- skill
      crps_clim <- Apply(list(data$obs$data), target_dims = time_dim,
                         fun = CRPS_clim, memb_dim = memb_dim,
                         clim.cross.val = cross.val)$output1
      crps_clim <- .drop_dims(crps_clim)
      skill_metrics[['crps_clim']] <- crps_clim
    # 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') {
      source("https://earth.bsc.es/gitlab/es/s2dv/-/raw/master/R/RMS.R")
      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,
                      ncores = ncores)
      } else {
        skill <- Bias(data$hcst$data, data$obs$data,
                      time_dim = time_dim,
                      memb_dim = memb_dim,
                      ncores = ncores)
      }
      skill <- .drop_dims(skill)
      skill_metrics[[ metric ]] <- skill
    # 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') {
      # Remove ensemble dim from obs to avoid veriApply warning
      obs_noensdim <- ClimProjDiags::Subset(data$obs$data, "ensemble", 1,
                                            drop = "selected")
      capture.output(
      skill <- easyVerification::veriApply(verifun = 'EnsSprErr',
                                           fcst = data$hcst$data,
                                           obs = obs_noensdim,
                                           tdim = which(names(dim(data$hcst$data))==time_dim),
                                           ensdim = which(names(dim(data$hcst$data))==memb_dim),
                                           na.rm = na.rm,
                                           ncpus = ncores)
      )
      remove(obs_noensdim)
      skill <- .drop_dims(skill)
      skill_metrics[[ metric ]] <- skill
    # 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
      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
    }
  }
  info(recipe$Run$logger, "##### SKILL METRIC COMPUTATION COMPLETE #####")
  .log_memory_usage(recipe$Run$logger, when = "After skill metric computation")
  # Save outputs
  #NURIA: I think change the output_dir is a problem for future savings
  if (recipe$Analysis$Workflow$Skill$save != 'none') {
    info(recipe$Run$logger, "##### 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, skill = skill_metrics,
                   data_cube = data$hcst, 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 = data$hcst)
      }
      # Save other skill metrics
      if (length(skill_metrics[-corr_metric_names]) > 0) {
        save_metrics(recipe = recipe, skill = skill_metrics[-corr_metric_names],
                     data_cube = data$hcst, agg = agg)
      }
    }
  }
 
  #Provenance
  if (recipe$Analysis$Provenance){
  if (is.null(graph.prov$parentnodename$indicesnodenames)){
    skill_metrics$graph.prov <-
      SunsetProv.Skill(recipe, graph.prov)
  }else {
    skill_metrics$graph.prov <-
      SunsetProv.Skill(recipe, graph.prov, Indices = TRUE)
  }
}

  #Return results
  return(skill_metrics)
}

Probabilities <- function(recipe, data) {
  ## 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, "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,
	 "##### PERCENTILES AND PROBABILITY CATEGORIES COMPUTED #####")
    .log_memory_usage(recipe$Run$logger, when = "After anomaly computation")
    # Save outputs
    if (recipe$Analysis$Workflow$Probabilities$save != 'none') {
      info(recipe$Run$logger,
           "##### 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)
  }
}

