
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 'bin' 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) {
  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 request more than only terciles:
  for  (ps in 1:length(exe_rps)) { ## Need to revise this loop, may not work for multiple categories!!
    if (tolower(recipe$Analysis$Horizon) == 'subseasonal') {
      data_crossval$probs$hcst_ev[[ps]] <- MergeDims(data_crossval$probs$hcst_ev[[ps]],
                                                     merge_dims = c('sweek', 'syear'),
                                                     rename_dim = 'syear', na.rm = FALSE)
      data_crossval$probs$obs_ev[[ps]] <- MergeDims(data_crossval$probs$obs_ev[[ps]],
                                                    merge_dims = c('sweek', 'syear'),
                                                    rename_dim = 'syear', na.rm = FALSE)
    }
  } ## close on ps
  
    if ('rps' %in% requested_metrics) {
      index <- which(sapply(categories, function(x) "1/3" %in% names(x)))
      rps <- RPS(exp = data_crossval$probs$hcst_ev[[index]],
                 obs = data_crossval$probs$obs_ev[[index]], memb_dim = NULL,
                 cat_dim = 'bin', cross.val = FALSE, time_dim = 'syear',
       	         Fair = fair, nmemb = nmemb,
                 ncores = ncores)
      rps_clim <- Apply(list(data_crossval$probs$obs_ev[[index]]),
                        target_dims = c('bin', 'syear'),
                        RPS_clim, bin_dim_abs = 'bin', 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) {
      index <- which(sapply(categories, function(x) "1/3" %in% names(x)))
      rps_syear <- RPS(exp = data_crossval$probs$hcst_ev[[index]],
                      obs = data_crossval$probs$obs_ev[[index]], memb_dim = NULL,
                      cat_dim = 'bin', 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) {
      index <- which(sapply(categories, function(x) "1/3" %in% names(x)))
      rps_clim_syear <- Apply(list(data_crossval$probs$obs_ev[[index]]),
                        target_dims = c('bin', 'syear'),
                        RPS_clim, bin_dim_abs = 'bin', 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) {
      index <- which(sapply(categories, function(x) "1/3" %in% names(x)))
      rpss <- RPSS(exp = data_crossval$probs$hcst_ev[[index]],
                   obs = data_crossval$probs$obs_ev[[index]],
                   ref = NULL, # ref is 1/3 by default if terciles
                   time_dim = 'syear', memb_dim = NULL,
                   cat_dim = 'bin', nmemb = nmemb,
                   dat_dim = NULL,
                   prob_thresholds = categories[[index]],
                   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)
      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 ('bss10' %in% requested_metrics) {
      index <- which(sapply(categories, function(x) "1/10" %in% names(x)))
      bss10 <- RPSS(exp = data_crossval$probs$hcst_ev[[index]],
                    obs = data_crossval$probs$obs_ev[[index]],
                    ref = NULL, # ref is 1/3 by default if terciles
                    time_dim = 'syear', memb_dim = NULL,
                    cat_dim = 'bin', nmemb = nmemb,
                    dat_dim = NULL,
                    prob_thresholds = 0.1,
                    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)
      skill_metrics$bss10 <- bss10$rpss
      skill_metrics$bss10_significance <- bss10$sign
    }
    if ('bss90' %in% requested_metrics) {
      index <- which(sapply(categories, function(x) "9/10" %in% names(x)))
      bss90 <- RPSS(exp = data_crossval$probs$hcst_ev[[index]],
                    obs = data_crossval$probs$obs_ev[[index]],
                    ref = NULL, # ref is 1/3 by default if terciles
                    time_dim = 'syear', memb_dim = NULL,
                    cat_dim = 'bin', nmemb = nmemb,
                    dat_dim = NULL,
                    prob_thresholds = 0.9,
                    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)
      skill_metrics$bss90 <- bss90$rpss
      skill_metrics$bss90_significance <- bss90$sign
    }
    

  if (tolower(recipe$Analysis$Horizon) == 'subseasonal') {
    # The evaluation of all metrics are done with extra sample
    data_crossval$hcst <- CST_MergeDims(data_crossval$hcst,
                                        merge_dims = c('sweek', 'syear'),
                                        rename_dim = 'syear', na.rm = FALSE)
    data_crossval$obs <- CST_MergeDims(data_crossval$obs,
                                       merge_dims = c('sweek', 'syear'),
                                       rename_dim = 'syear', na.rm = FALSE)
    data_crossval$ref_obs_tr <- MergeDims(data_crossval$ref_obs_tr,
                                          merge_dims = c('sweek', 'syear'),
                                          rename_dim = 'syear', na.rm = FALSE)
    data_crossval$hcst.full_val <- CST_MergeDims(data_crossval$hcst.full_val,
                                                 merge_dims = c('sweek', 'syear'),
                                                 rename_dim = 'syear', na.rm = FALSE)
    data_crossval$obs.full_val <- CST_MergeDims(data_crossval$obs.full_val,
                                                merge_dims = c('sweek', 'syear'),
                                                rename_dim = 'syear', na.rm = FALSE)
  }
  
  
  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) {
    ## Temporarily changing input data to show mean bias of calibrated data
    # if (!is.null(data_crossval$hcst$data)) {
      mean_bias <- Bias(exp = data_crossval$hcst$data,
                        obs = data_crossval$obs$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
  }
  # 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,
  #          "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
  }
  original <- recipe$Run$output_dir
  recipe$Run$output_dir <- paste0(original, "/outputs/Skill/")

  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)
		      })
  # Save metrics
  if (tolower(recipe$Analysis$Workflow$Skill$save) == 'all') {
    save_metrics(recipe = recipe,
                 metrics = skill_metrics,
                 data_cube = data_crossval$hcst, agg = 'global',
                 outdir = recipe$Run$output_dir)
  }
  
  recipe$Run$output_dir <- original
  # reduce dimension to work with Visualization module:
  skill_metrics <- lapply(skill_metrics, function(x) {drop(x)})
  skill_metrics <- lapply(skill_metrics, function(x){
                          InsertDim(x, pos = 1, len = 1, name = 'var')})
  return(skill_metrics)
}

