# Full-cross-val workflow
## This code should be valid for individual months and temporal averages
source("modules/Crossval/R/tmp/GetProbs.R")
  
Crossval_anomalies <- function(recipe, data) {
  cross.method <- recipe$Analysis$cross.method
  # TODO move check
  if (is.null(cross.method)) {
    cross.method <- 'leave-one-out'
  }
  categories <- recipe$Analysis$Workflow$Probabilities$percentiles
  categories <- lapply(categories, function (x) { 
                       sapply(x, function(y) {
                              eval(parse(text = y))})})
  ncores <- recipe$Analysis$ncores
  na.rm <- recipe$Analysis$remove_NAs 
  ## data dimensions 
  sdate_dim <- dim(data$hcst$data)['syear']
  orig_dims <- names(dim(data$hcst$data))
  # spatial dims
  if ('latitude' %in% names(dim(data$hcst$data))) {
    nlats <- dim(data$hcst$data)['latitude']
    nlons <- dim(data$hcst$data)['longitude']
    agg = 'global'
  } else if  ('region' %in% names(dim(data$hcst$data))) {
    agg = 'region'
    nregions <- dim(data$hcst$data)['region']
  }
  # TODO fix it to use new version https://earth.bsc.es/gitlab/external/cstools/-/blob/dev-cross-indices/R/CST_Calibration.R#L570
  cross <- CSTools:::.make.eval.train.dexes(eval.method = cross.method,
                                            amt.points = sdate_dim,
                                            amt.points_cor = NULL) # k = ?

  # Define function to compute anomalies in full cross-validation
  .cross_anomalies <- function(indices, exp, obs, cross, categories, recipe) {
    ## Objects to return:
    # lims_ano_hcst_tr: hcst percentiles (per category)
    # lims_ano_obs_tr: obs percentiles (per category)
    # ano_hcst_ev: hcst cross-validated anomalies
    # ano_obs_ev: obs cross-validated anomalies
    # ano_obs_tr: obs training anomalies 
    
    # subset years: Subset works at BSC not at Athos
    ## training indices
    cross <- cross[[indices]]
    print(paste("Crossval Index:", cross$eval.dexes))
    obs_tr <- Subset(obs, along = 'syear',
   		     indices = cross$train.dexes)
    hcst_tr <- Subset(exp, along = 'syear', 
  		      indices = cross$train.dexes)
    ## evaluation indices
    hcst_ev <- Subset(exp, along = 'syear',
  		      indices = cross$eval.dexes, drop = 'selected')
    obs_ev <- Subset(obs, along = 'syear',
  		     indices = cross$eval.dexes, drop = 'selected')
    if (tolower(recipe$Analysis$Horizon) == 'subseasonal') {
      central_day <- (dim(exp)['sday'] + 1)/2
      hcst_tr <- MergeDims(hcst_tr, merge_dims = c('sday', 'syear'), 
                          rename_dim = 'syear', na.rm = FALSE) 

      obs_tr <- MergeDims(obs_tr, merge_dims = c('sday', 'syear'),
                          rename_dim = 'syear', na.rm = FALSE)
      # 'sday' dim to select the central day
      hcst_ev <- Subset(hcst_ev, along = 'sday', indices = central_day) 

    }
    # compute climatology:
    clim_obs_tr <- MeanDims(obs_tr, 'syear')
    clim_hcst_tr <- MeanDims(hcst_tr, c('syear', 'ensemble'))
    # compute anomalies:
    ano_obs_tr <- s2dv::Ano(obs_tr, clim_obs_tr,
                            ncores = recipe$Analysis$ncores)
    ano_hcst_tr <- s2dv::Ano(hcst_tr, clim_hcst_tr,
                             ncores = recipe$Analysis$ncores)
    ano_hcst_ev <- s2dv::Ano(hcst_ev, clim_hcst_tr,
                             ncores = recipe$Analysis$ncores)
    ano_obs_ev <- s2dv::Ano(obs_ev, clim_obs_tr,
                            ncores = recipe$Analysis$ncores)
    # compute category limits
    lims_ano_hcst_tr <- Apply(ano_hcst_tr, target_dims = c('syear', 'ensemble'),
                              fun = function(x, prob_lims) {
                                lapply(prob_lims, function(ps) {
                                  as.array(quantile(as.vector(x), ps,
                                                    na.rm = na.rm))})},
                              prob_lims = list(unlist(categories)),
                              ncores = ncores)$output1
    lims_ano_obs_tr <- Apply(ano_obs_tr, target_dims = c('syear', 'ensemble'),
                             fun = function(x, prob_lims) {
                               lapply(prob_lims, function(ps) {
                                 as.array(quantile(as.vector(x), ps,
                                                   na.rm = na.rm))})},
                             prob_lims = list(unlist(categories)),
                             ncores = ncores)$output1
    if (tolower(recipe$Analysis$Horizon) == 'subseasonal') {
     ano_hcst_tr <- SplitDim(ano_hcst_tr, split_dim = 'syear',
                             new_dim_name = 'sday',
                             indices = rep(1:dim(data$hcst$data)['sday'],
                                           length(cross[[t]]$train.dexes)))
     ano_obs_tr <- SplitDim(obs_tr, split_dim = 'syear', new_dim_name = 'sday',
                            indices = rep(1:dim(data$hcst$data)['sday'],
                                          length(cross[[t]]$train.dexes)))
    }

    return(list(ano_hcst_ev = ano_hcst_ev,
                ano_obs_ev = ano_obs_ev,
                ano_obs_tr = ano_obs_tr,
                lims_ano_hcst_tr = lims_ano_hcst_tr,
                lims_ano_obs_tr = lims_ano_obs_tr))
  }

  output_dims <- names(dim(data$hcst$data))[which(names(dim(data$hcst$data)) != "syear")]
  lims_output_dims <- c("bin",
                        names(dim(data$hcst$data))[which(!names(dim(data$hcst$data)) %in% c("syear", "ensemble"))])
  indices <- array(1:length(cross), dim = c(syear = length(cross)))
  res <- Apply(data = list(indices = indices,
                           exp = data$hcst$data,
                           obs = data$obs$data),
               target_dims = list(indices = NULL,
                                  obs = names(dim(data$obs$data)),
                                  exp = names(dim(data$hcst$data))),
               output_dims = list(ano_hcst_ev = output_dims,
                                  ano_obs_ev = output_dims,
                                  ano_obs_tr = c("dat", "var", "sday", "sweek",
                                                 "ensemble", "time", "latitude", "longitude",
                                                 "unneeded"),
                                  lims_ano_hcst_tr = lims_output_dims,
                                  lims_ano_obs_tr = lims_output_dims),
               fun = .cross_anomalies,
               ncores = 1, ## TODO: Explore options for core distribution
               cross = cross,
               categories = categories,
               recipe = recipe)

  info(recipe$Run$logger,
       "#### Anomalies Cross-validation loop ended #####")

  # Rearrange limits into a list:
  lims_tmp_hcst <- list()
  lims_tmp_obs <- list()
  last_index <- 0
  for (ps in 1:length(categories)) {
    indices <- last_index + 1:length(categories[[ps]])
    lims_category_hcst <- Subset(res$lims_ano_hcst_tr, along = "bin",
                                 indices = list(indices),
                                 drop = FALSE)
    lims_tmp_hcst[[ps]] <- lims_category_hcst
    lims_category_obs <- Subset(res$lims_ano_obs_tr, along = "bin",
                                indices = list(indices),
                                drop = FALSE)
    lims_tmp_obs[[ps]] <- lims_category_obs
    last_index <- indices[length(indices)]
  }
  res$lims_ano_hcst_tr <- lims_tmp_hcst
  res$lims_ano_obs_tr <- lims_tmp_obs
  rm(lims_tmp_hcst, lims_tmp_obs)
  gc()
  
  # To make crps_clim to work the reference forecast need to have same dims as obs:
  res$ano_obs_tr <- Subset(res$ano_obs_tr, along = 'unneeded',
                           indices = 1, drop = 'selected')

  recipe$Run$output_dir <- paste0(recipe$Run$output_dir,
                                  "/outputs/Anomalies/")

  # Forecast anomalies:
  if (!is.null(data$fcst)) {
    clim_hcst <- Apply(data$hcst$data,
                       target_dims = c('syear', 'ensemble'),
                       mean,
                       na.rm = na.rm,
                       ncores = ncores)$output1
    data$fcst$data <- Ano(data = data$fcst$data, clim = clim_hcst)
    hcst_ano <- Ano(data = data$hcst$data, clim = clim_hcst)
    # Terciles limits using the whole hindcast period:
    print(categories)
    lims_fcst <- Apply(hcst_ano, target_dims = c('syear', 'ensemble'),
                       fun = function(x, prob_lims) {
                         lapply(prob_lims, function(ps) {
                                as.array(quantile(as.vector(x), ps,
                                                  na.rm = na.rm))})},
                       output_dims = lapply(categories, function(x) {'bin'}),
                       prob_lims = categories,
                       ncores = ncores)
    clim_obs <- Apply(data$obs$data,
                       target_dims = c('syear', 'ensemble'),
                       mean,
                       na.rm = na.rm,
                       ncores = ncores)$output1
    obs_ano <- Ano(data = data$obs$data, clim = clim_obs)
    lims <- Apply(obs_ano, target_dims = c('syear', 'ensemble'),
                  fun = function(x, prob_lims) {
                    lapply(prob_lims, function(ps) {
                             as.array(quantile(as.vector(x), ps,
                                               na.rm = na.rm))})},
                  output_dims = lapply(categories, function(x) {'bin'}),
                  prob_lims = categories,
                  ncores = ncores)
    tmp_lims2 <- list()
    for (ps in 1:length(categories)) {
      tmp_lims3 <- .drop_dims(lims[[ps]])
      for (l in 1:dim(lims[[ps]])['bin']) {
        tmp_lims <- ClimProjDiags::Subset(x = tmp_lims3,
                                          along = "bin",
                                          indices = l)
        tmp_lims2 <- append(tmp_lims2, list(tmp_lims))
        names(tmp_lims2)[length(tmp_lims2)] <- as.character(categories[[ps]][l])
      }
      if (recipe$Analysis$Workflow$Probabilities$save == 'yes') {
        save_percentiles(recipe = recipe, percentiles = tmp_lims2,
                         data_cube = data$obs,
                         agg = "global", outdir = NULL)
      }
    }
  }

  ## TODO: Should this happen here?
  # Make categories rounded number to use as names:
  categories <- lapply(categories, round, digits = 2)

  # Compute Probabilities
  fcst_probs <- lapply(categories, function(x){NULL})
  hcst_probs_ev <- lapply(categories, function(x){NULL})
  obs_probs_ev <- lapply(categories, function(x){NULL})

  for (ps in 1:length(categories)) {
    hcst_probs_ev[[ps]] <- GetProbs(res$ano_hcst_ev, time_dim = 'syear',
                                    prob_thresholds = NULL,
                                    bin_dim_abs = 'bin',
                                    indices_for_quantiles = NULL,
                                    memb_dim = 'ensemble',
                                    abs_thresholds = res$lims_ano_hcst_tr[[ps]],
                                    ncores = ncores)
    obs_probs_ev[[ps]] <- GetProbs(res$ano_obs_ev, time_dim = 'syear',
                                   prob_thresholds = NULL,
                                   bin_dim_abs = 'bin',
                                   indices_for_quantiles = NULL,
                                   memb_dim = 'ensemble',
                                   abs_thresholds = res$lims_ano_obs_tr[[ps]],
                                   ncores = ncores)
    if (!is.null(data$fcst)) {
      fcst_probs[[ps]] <- GetProbs(data$fcst$data, time_dim = 'syear',
                                   prob_thresholds = NULL,
                                   bin_dim_abs = 'bin',
                                   indices_for_quantiles = NULL,
                                   memb_dim = 'ensemble',
                                   abs_thresholds = lims_fcst[[ps]],
                                   ncores = ncores)
    }
  }
  # Convert to s2dv_cubes the resulting anomalies
  ano_hcst <- data$hcst
  ano_hcst$data <- res$ano_hcst_ev
  ano_obs <- data$obs
  ano_obs$data <- res$ano_obs_ev

  info(recipe$Run$logger,
       "#### Anomalies and Probabilities Done #####")
  if (recipe$Analysis$Workflow$Anomalies$save != 'none') {
    info(recipe$Run$logger, "##### START SAVING ANOMALIES #####")
    recipe$Run$output_dir <- paste0(recipe$Run$output_dir,
                                    "/outputs/Anomalies/")
    # Save forecast
    if ((recipe$Analysis$Workflow$Anomalies$save %in%
         c('all', 'exp_only', 'fcst_only')) && !is.null(data$fcst)) {
      save_forecast(recipe = recipe, data_cube = data$fcst, type = 'fcst')
    }
    # Save hindcast
    if (recipe$Analysis$Workflow$Anomalies$save %in%
        c('all', 'exp_only')) {
      save_forecast(recipe = recipe, data_cube = ano_hcst, type = 'hcst')
    }
    # Save observation
    if (recipe$Analysis$Workflow$Anomalies$save == 'all') {
      save_observations(recipe = recipe, data_cube = ano_obs)
    }
  }
  # Save probability bins
  probs_hcst <- list()
  probs_fcst <- list()
  probs_obs <- list()
  all_names <- NULL

  for (ps in 1:length(categories)) {
    for (perc in 1:(length(categories[[ps]]) + 1)) {
      if (perc == 1) {
        name_elem <- paste0("below_", categories[[ps]][perc])
      } else if (perc == length(categories[[ps]]) + 1) {
        name_elem <- paste0("above_", categories[[ps]][perc-1])
      } else {
        name_elem <- paste0("from_", categories[[ps]][perc-1],
                            "_to_", categories[[ps]][perc])
      }
      probs_hcst <- append(probs_hcst,
                           list(Subset(hcst_probs_ev[[ps]],
                                       along = 'bin',
                                       indices = perc,
                                       drop = 'selected')))
      probs_obs <- append(probs_obs,
                          list(Subset(obs_probs_ev[[ps]],
                                      along = 'bin',
                                      indices = perc,
                                      drop = 'selected')))
      if (!is.null(data$fcst)) {
        probs_fcst <- append(probs_fcst,
                             list(Subset(fcst_probs[[ps]],
                                         along = 'bin',
                                         indices = perc,
                                         drop = 'selected')))
      }
      all_names <- c(all_names, name_elem)
    }
  }
  ## TODO: Apply .drop_dims() directly to hcst_probs_ev etc; and move
  ## reorganizing and renaming to save_probabilities()
  names(probs_hcst) <- all_names
  probs_hcst <- lapply(probs_hcst, .drop_dims)
  names(probs_obs) <- all_names
  probs_obs <- lapply(probs_obs, .drop_dims)
  if (!is.null(data$fcst)) {
    names(probs_fcst) <- all_names
    probs_fcst <- lapply(probs_fcst, .drop_dims)
  }

  if (recipe$Analysis$Workflow$Probabilities$save %in%
    c('all', 'bins_only')) {
    save_probabilities(recipe = recipe, probs = probs_hcst,
                       data_cube = data$hcst, agg = agg,
                       type = "hcst")
    save_probabilities(recipe = recipe, probs = probs_obs,
                       data_cube = data$hcst, agg = agg,
                       type = "obs")
    if (!is.null(probs_fcst)) {
      save_probabilities(recipe = recipe, probs = probs_fcst,
                         data_cube = data$fcst, agg = agg,
                         type = "fcst")
    }
  }
  # Save ensemble mean for multimodel option:
  hcst_EM <- MeanDims(ano_hcst$data, 'ensemble', drop = T)
#  save_metrics(recipe = recipe,
#               metrics = list(hcst_EM =
#                 Subset(hcst_EM, along = 'dat', indices = 1, drop = 'selected')),
#               data_cube = data$hcst, agg = agg,
#               module = "statistics")
#  fcst_EM <- NULL
  if (!is.null(data$fcst)) {
    fcst_EM <- MeanDims(data$fcst$data, 'ensemble', drop = T)
#    save_metrics(recipe = recipe,
#                 metrics = list(fcst_EM =
#                   Subset(fcst_EM, along = 'dat', indices = 1, drop = 'selected')),
#                 data_cube = data$fcst, agg = agg,
#                 module = "statistics")
  }
  return(list(hcst = ano_hcst, obs = ano_obs, fcst = data$fcst,
              hcst.full_val = data$hcst, obs.full_val = data$obs,
              hcst_EM = hcst_EM, fcst_EM = fcst_EM,
              #cat_lims = list(hcst_tr = res$lims_ano_hcst_tr,
              #                obs_tr = res$lims_ano_obs_tr),
              probs = list(hcst_ev = hcst_probs_ev,
                           obs_ev = obs_probs_ev,
			   fcst = fcst_probs),
              ref_obs_tr = res$ano_obs_tr))
}


## The result contains the inputs for Skill_full_crossval.
## 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

