
source("modules/Saving/Saving.R")
source("modules/Crossval/R/tmp/RPS.R")
source("modules/Crossval/R/RPS_clim.R")
source("modules/Crossval/R/CRPS_clim.R")
source("modules/Crossval/R/tmp/RPSS.R")
source("modules/Crossval/R/tmp/RandomWalkTest.R")
source("modules/Crossval/R/tmp/Corr.R")
source("modules/Crossval/R/tmp/Bias.R")
source("modules/Crossval/R/tmp/SprErr.R")
source("modules/Crossval/R/tmp/Eno.R")
source("modules/Skill/R/tmp/CRPS.R")

## data_crossval is the result from function full_crossval_anomalies or similar.
## this is a list with the required elements:
  ## probs is a list with
    ## probs$hcst_ev and probs$obs_ev 
    ## probs$hcst_ev will have as many elements in the $Probabilities$percentiles
    ## each element will be an array with 'cat' dimension
    ## the same for probs$obs_ev
  ## hcst is a s2dv_cube for the post-processed hindcast for the evalutaion indices
    ## in this case cross validated anomalies
  ## obs is a s2dv_cube for the post-processed obs 
    ## in this case cross validated anomalies
  ## fcst is a s2dv_cube for the post-processed fcst
    ## in this case cross anomalies with the full hindcast period
    ## this object is not required for skill assessment
  ## hcst.full_val and obs.full_val are the original data to compute mean bias
  ## cat_lims used to compute the probabilities
    ## this object is not required for skill assessment
  ## ref_obs_tr is an array with the cross-validate observed anomalies
    ## to be used as reference forecast in the CRPSS and CRPS_clim
    ## it is computed from the training indices
## the recipe could be used to read the Percentiles 
## if fair is TRUE, the nmemb used to compute the probabilities is needed
  ## nmemb_ref is the number of year - 1 in case climatological forecast is the reference
Crossval_metrics <- function(recipe, data_crossval,
                             fair = FALSE, nmemb = NULL, nmemb_ref = NULL,
                             retrieve = TRUE, nchunks = nchunks) {
  ncores <- recipe$Analysis$ncores
  alpha <- recipe$Analysis$Skill$alpha
  na.rm <- recipe$Analysis$remove_NAs
  categories <- recipe$Analysis$Workflow$Probabilities$percentiles
  categories <- lapply(categories, function (x) { 
                       sapply(x, function(y) {
                              eval(parse(text = y))})})

  # TODO: distinguish between rpss and bss 
  # if 1 percentile -> bss
  # if more than 1 -> rpss 
  exe_rps <- unlist(lapply(categories, function(x) {
                    if (length(x) > 1) {
                      x <- x[1] *100
                      }
                      return(x)}))
  if (is.null(alpha)) {
    alpha <- 0.05
  }
  ## START SKILL ASSESSMENT:
  skill_metrics <- list()
  requested_metrics <- strsplit(recipe$Analysis$Workflow$Skill$metric,
                                ", | |,")[[1]]
  # The recipe allows to requset more than only terciles:
  for (ps in 1:length(exe_rps)) {
    if ('rps' %in% requested_metrics) {
      rps <- RPS(exp = data_crossval$probs$hcst_ev[[ps]],
                 obs = data_crossval$probs$obs_ev[[1]], memb_dim = NULL,
                 cat_dim = 'cat', cross.val = FALSE, time_dim = 'syear',
                 Fair = fair, nmemb = nmemb,
                 ncores = ncores)
      rps_clim <- Apply(list(data_crossval$probs$obs_ev[[1]]),
                        target_dims = c('cat', 'syear'),
                        RPS_clim, bin_dim_abs = 'cat', Fair = fair,
      	  	        cross.val = FALSE, ncores = ncores)$output1
      skill_metrics$rps <- rps
      skill_metrics$rps_clim <- rps_clim
      # names based on the categories:
      # To use it when visualization works for more rps
      #skill_metrics[[paste0('rps', exe_rps[ps])]] <- rps
      #skill_metrics[[paste0('rps_clim',
      #                      exe_rps[ps])]] <- rps_clim 
    } 
    if ('rps_syear' %in% requested_metrics) {
      rps_syear <- RPS(exp = data_crossval$probs$hcst_ev[[ps]],
                      obs = data_crossval$probs$obs_ev[[1]], memb_dim = NULL,
                      cat_dim = 'cat', cross.val = FALSE, time_dim = 'syear',
                      Fair = fair, nmemb = nmemb, return_mean = FALSE,
                      ncores = ncores)
      skill_metrics$rps_syear <- rps_syear
    }
    if ('rps_clim_syear' %in% requested_metrics) {
      rps_clim_syear <- Apply(list(data_crossval$probs$obs_ev[[1]]),
                        target_dims = c('cat', 'syear'),
                        RPS_clim, bin_dim_abs = 'cat', Fair = fair,
                        cross.val = FALSE, return_mean = FALSE, 
                        output_dims = 'syear', ncores = ncores)$output1
      skill_metrics$rps_clim_syear <- rps_clim_syear
    }
    if ('rpss' %in% requested_metrics) {
      rpss <- RPSS(exp = data_crossval$probs$hcst_ev[[1]],
                   obs = data_crossval$probs$obs_ev[[1]],
                   ref = NULL, # ref is 1/3 by default if terciles
                   time_dim = 'syear', memb_dim = NULL,
                   cat_dim = 'cat', nmemb = nmemb,
                   dat_dim = NULL,
                   prob_thresholds = categories[[ps]],
                   indices_for_clim = NULL,
                   Fair = fair, weights_exp = NULL, weights_ref = NULL, 
                   cross.val = FALSE, na.rm = na.rm, 
                   sig_method.type = 'two.sided.approx', alpha = alpha,
                   ncores = ncores)
      ## TODO: These lines are probably not needed
      skill_metrics$rpss <- rpss$rpss
      skill_metrics$rpss_significance <- rpss$sign
      # TO USE IT when visualization works for more rpsss
      #skill_metrics[[paste0('rpss', exe_rps[ps])]] <- rpss$rpss
      #skill_metrics[[paste0('rpss',
      #                      exe_rps[ps], 
      #                      "_significance")]] <- rpss$sign
    }
  }
  if ('crps' %in% requested_metrics) {
    crps <- CRPS(exp  = data_crossval$hcst$data,
                 obs = data_crossval$obs$data,
                 time_dim = 'syear', memb_dim = 'ensemble',
                 Fair = fair,
                 ncores = ncores)
    skill_metrics$crps <- crps
    crps_clim <- CRPS(exp = data_crossval$ref_obs_tr,
                      obs = data_crossval$obs$data,
                      time_dim = 'syear', memb_dim = 'ensemble',
                      Fair = fair, ncores = ncores) 
    skill_metrics$crps_clim <- crps_clim
  }
  if ('crps_syear' %in% requested_metrics) {
    crps_syear <- CRPS(exp  = data_crossval$hcst$data,
                       obs = data_crossval$obs$data,
                       time_dim = 'syear', memb_dim = 'ensemble',
                       Fair = fair, return_mean = FALSE, 
                       ncores = ncores)
    skill_metrics$crps_syear <- crps_syear
  }
  if ('crps_clim_syear' %in% requested_metrics) {
    crps_clim_syear <- CRPS(exp = data_crossval$ref_obs_tr,
                            obs = data_crossval$obs$data,
                            time_dim = 'syear', memb_dim = 'ensemble',
                            Fair = fair, return_mean = FALSE, ncores = ncores) 
    skill_metrics$crps_clim_syear <- crps_clim_syear
  }
  if ('crpss' %in% requested_metrics) {
    crpss <- CRPSS(exp = data_crossval$hcst$data,
                   obs = data_crossval$obs$data,
                   ref = data_crossval$ref_obs_tr,
                   memb_dim = 'ensemble', Fair = fair,
                   time_dim = 'syear', clim.cross.val = FALSE,
                   ncores = ncores)
    skill_metrics$crpss <- crpss$crpss
    skill_metrics$crpss_significance <- crpss$sign
  }

  if ('enscorr' %in% requested_metrics) {
    enscorr <- Corr(exp = data_crossval$hcst$data,
                    obs = data_crossval$obs$data, 
                    dat_dim = NULL,
                    time_dim = 'syear', 
                    method = 'pearson', 
                    memb_dim = 'ensemble',
                    memb = F,
                    conf = F,
                    pval = F,
                    sign = T,
                    alpha = alpha,
                    ncores = ncores)
    skill_metrics$enscorr <- enscorr$corr
    skill_metrics$enscorr_significance <- enscorr$sign 
  }
  if ('mean_bias' %in% requested_metrics) {
    if (!is.null(data_crossval$hcst.full_val$data)) {
      mean_bias <- Bias(exp = data_crossval$hcst.full_val$data,
                        obs = data_crossval$obs.full_val$data,
                        time_dim = 'syear',
                        memb_dim = 'ensemble',
                        alpha = alpha,
                        ncores = ncores)
      skill_metrics$mean_bias <- mean_bias$bias
      skill_metrics$mean_bias_significance <- mean_bias$sig
    } else {
      info(recipe$Run$logger, retrieve = retrieve,
           "Full values not available")
    }
  }
  if ('enssprerr' %in% requested_metrics) {
    enssprerr <- SprErr(exp = data_crossval$hcst$data,
                        obs = data_crossval$obs$data,
                        memb_dim = 'ensemble', dat_dim = NULL,
                        time_dim = 'syear', pval = TRUE,
                        ncores = ncores)
    skill_metrics$enssprerr <- enssprerr$ratio
    skill_metrics$enssprerr_significance <- enssprerr$p.val <= alpha
  }
  if ('rms' %in% requested_metrics) {
    rms <- RMS(exp = data_crossval$hcst$data,
               obs = data_crossval$obs$data,
               memb_dim = 'ensemble', dat_dim = NULL,
               time_dim = 'syear', alpha = alpha,
               ncores = ncores)
    skill_metrics$rms <- rms$rms
  }
  if ('rmss' %in% requested_metrics) {
    rmss <- RMSSS(exp = data_crossval$hcst$data,
                  obs = data_crossval$obs$data,
                  ref = res$ref_obs_tr,
                  memb_dim = 'ensemble', dat_dim = NULL,
                  time_dim = 'syear', alpha = alpha, sign = TRUE,
                  ncores = ncores)
    skill_metrics$rmss <- rmss$rmss
    skill_metrics$rmss_significance <- rmss$sign
  }
  if (is.null(data_crossval$hcst_EM)) {
    data_crossval$hcst_EM <- MeanDims(data_crossval$hcst$data,
                                      dims = 'ensemble',
                                      drop = TRUE)
  }
  if (any(c('std', 'standard_deviation') %in% requested_metrics)) {
      std_hcst <- Apply(data = data_crossval$hcst_EM, 
                        target_dims = 'syear', 
                        fun = 'sd')$output1
      
      std_obs <- Apply(data = data_crossval$obs$data, 
                       target_dims = 'syear', 
                       fun = 'sd')$output1
      
      skill_metrics[['std_hcst']] <- std_hcst
      skill_metrics[['std_obs']] <- std_obs
  }
  if (any(c('var', 'variance') %in% requested_metrics)) {
    ## Calculate variance
    var_hcst <- Apply(data = data_crossval$hcst_EM, 
                      target_dims = 'syear', 
                      fun = 'sd')$output1 ^ 2
      
    var_obs <- Apply(data = data_crossval$obs$data, 
                     target_dims = 'syear', 
                     fun = 'sd')$output1 ^ 2
      
    skill_metrics[['var_hcst']] <- var_hcst
    skill_metrics[['var_obs']] <- var_obs
  } ## close if on variance
  if ('n_eff' %in% requested_metrics) {   
    ## Calculate degrees of freedom
    n_eff <- s2dv::Eno(data = data_crossval$obs$data,
                       time_dim = 'syear',
                       na.action = na.pass,
                       ncores = ncores)
    skill_metrics[['n_eff']] <- n_eff
  } ## close on n_eff  

  if (any(c('cov', 'covariance') %in% requested_metrics)) {
    covariance <- Apply(data = list(x = data_crossval$obs$data,
                                    y = data_crossval$hcst_EM), 
                         target_dims = 'syear',
                         fun = function(x, y) {
                           cov(as.vector(x), as.vector(y),
                               use = "everything", 
                               method = "pearson")})$output1
   skill_metrics$covariance <- covariance
  }

  skill_metrics <- lapply(skill_metrics,
                          function(x) {
                            if (is.logical(x)) {
                              dims <- dim(x)
                              res <- as.numeric(x)
                              dim(res) <- dims
                            } else {
                              res <- x
                            }
                            return(res)
                          })

  # reduce dimension to work with Visualization module:
  skill_metrics <- lapply(skill_metrics, function(x) {.drop_dims(x)})

  ## TODO: Generalize this or find a better way to handle it
  # Separate skill metrics that have syear dimension
  syear_metrics <- grepl("syear", names(skill_metrics))
  skill_metrics_syear <- skill_metrics[syear_metrics]
  skill_metrics <- skill_metrics[!syear_metrics]

  # Save metrics
  if (retrieve) {
    ## TODO: Update this part
    recipe$Run$output_dir <- paste0(recipe$Run$output_dir, "/outputs/Skill/")
    save_metrics(recipe = recipe,
                 metrics = skill_metrics,
                 data_cube = data_crossval$hcst,
                 agg = 'global',
                 outdir = recipe$Run$output_dir)
  } else {
    source("tools/save_metadata_chunks.R")
    if (length(skill_metrics) > 0) {
      skill_metrics <- save_metadata_chunks(recipe = recipe,
                                            metrics = skill_metrics,
                                            data_cube = data_crossval$hcst,
                                            module = "Skill",
                                            nchunks = nchunks)
    } else {
      skill_metrics <- NULL
    }
    if (length(skill_metrics_syear) > 0) {
      skill_metrics_syear <- save_metadata_chunks(recipe = recipe,
                                                  metrics = skill_metrics_syear,
                                                  data_cube = data_crossval$hcst,
                                                  module = "Skill_syear",
                                                  nchunks = nchunks)
    } else {
      skill_metrics_syear <- NULL
    }
    skill_metrics <- list(skill = skill_metrics,
                          skill_syear = skill_metrics_syear)
  }
  return(skill_metrics)
}

