# 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']
  }
  # output_dims from loop base on original dimensions
  ## ex: 'dat', 'var', 'sday', 'sweek', 'ensemble', 'time', 
  ##     'latitude', 'longitude', 'unneeded', 'syear'
  ev_dim_names <-  c(orig_dims[-which(orig_dims %in% 'syear')],
                     names(sdate_dim))
  orig_dims[orig_dims %in% 'ensemble'] <- 'unneeded'
  orig_dims[orig_dims %in% 'syear'] <- 'ensemble'
  tr_dim_names <-c(orig_dims,
                     names(sdate_dim))
  # 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 = ?

  ## output objects
  ano_hcst_ev_res <- NULL
  ano_obs_ev_res <- NULL
  ano_obs_tr_res <- NULL
  # as long as probs requested in recipe:
  lims_ano_hcst_tr_res <- lapply(categories, function(X) {NULL})
  lims_ano_obs_tr_res <- lapply(categories, function(X) {NULL})

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

  for (t in 1:length(cross)) {
    info(recipe$Run$logger, paste("crossval:", t))

    # subset years: Subset works at BSC not at Athos
    ## training indices
    obs_tr <- Subset(data$obs$data, along = 'syear',
   		     indices = cross[[t]]$train.dexes)
    hcst_tr <- Subset(data$hcst$data, along = 'syear', 
  		      indices = cross[[t]]$train.dexes)
    ## evaluation indices
    hcst_ev <- Subset(data$hcst$data, along = 'syear',
  		      indices = cross[[t]]$eval.dexes, drop = 'selected')
    obs_ev <- Subset(data$obs$data, along = 'syear',
  		     indices = cross[[t]]$eval.dexes, drop = 'selected')
    # 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 = ncores)
    ano_hcst_tr <- s2dv::Ano(hcst_tr, clim_hcst_tr,
                             ncores = ncores)
    ano_hcst_ev <- s2dv::Ano(hcst_ev, clim_hcst_tr,
                             ncores = ncores)
    ano_obs_ev <- s2dv::Ano(obs_ev, clim_obs_tr,
                            ncores = 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) {
                                  quantile(as.vector(x), ps, na.rm = na.rm)})},
                              output_dims = lapply(categories, function(x) {'bin'}), 
                              prob_lims = categories,
                              ncores = ncores)
    lims_ano_obs_tr <- Apply(ano_obs_tr, target_dims = c('syear'),#, 'ensemble'),
                             fun = function(x, prob_lims) {
                               lapply(prob_lims, function(ps) {
                                 quantile(as.vector(x), ps, na.rm = na.rm)})},
                             output_dims = lapply(categories, function(x){'bin'}),
                             prob_lims = categories,
                             ncores = ncores)
    #store results
    ano_hcst_ev_res <- abind(ano_hcst_ev_res, ano_hcst_ev,
                             along = length(dim(ano_hcst_ev)) + 1)
    ano_obs_ev_res <- abind(ano_obs_ev_res, ano_obs_ev,
                            along = length(dim(ano_obs_ev)) + 1)
    ano_obs_tr_res <- abind(ano_obs_tr_res, ano_obs_tr,
                            along = length(dim(ano_obs_tr)) + 1)
    for(ps in 1:length(categories)) {
      lims_ano_hcst_tr_res[[ps]] <- abind(lims_ano_hcst_tr_res[[ps]], lims_ano_hcst_tr[[ps]],
                                along = length(dim(lims_ano_hcst_tr[[ps]])) + 1)
      lims_ano_obs_tr_res[[ps]] <- abind(lims_ano_obs_tr_res[[ps]], lims_ano_obs_tr[[ps]],
                                 along = length(dim(lims_ano_obs_tr[[ps]])) + 1)
    }
  }
  info(recipe$Run$logger,
       "#### Anomalies Cross-validation loop ended #####")
  gc()
  # Add dim names:
  names(dim(ano_hcst_ev_res)) <- ev_dim_names
  names(dim(ano_obs_ev_res)) <- ev_dim_names
  names(dim(ano_obs_tr_res)) <- tr_dim_names
  # To make crps_clim to work the reference forecast need to have same dims as obs:
  ano_obs_tr_res <- Subset(ano_obs_tr_res, along = 'unneeded',
                           indices = 1, drop = 'selected')
  for(ps in 1:length(categories)) {
    names(dim(lims_ano_hcst_tr_res[[ps]])) <- c('bin',
          orig_dims[-which(orig_dims %in% c('ensemble', 'unneeded'))], 'syear')
    names(dim(lims_ano_obs_tr_res[[ps]])) <- c('bin',
          tr_dim_names[-which(tr_dim_names %in% c('ensemble'))])
    lims_ano_obs_tr_res[[ps]] <- Subset(lims_ano_obs_tr_res[[ps]],
                                        along = 'unneeded', indices = 1, drop = 'selected')
  }
  # Make categories rounded number to use as names:
  categories <- recipe$Analysis$Workflow$Probabilities$percentiles
  categories <- lapply(categories, function (x) {
                       sapply(x, function(y) {
                              round(eval(parse(text = y)),2)})})
  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:
    lims_fcst <- Apply(hcst_ano, target_dims = c('syear', 'ensemble'),
                       fun = function(x, prob_lims) {
                         lapply(prob_lims, function(ps) {
                         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) {
                         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(lims[[ps]])
      for (l in 1:dim(lims[[ps]])['bin']) {
        tmp_lims <- tmp_lims3[l,,,]
        if (!('var' %in% names(dim(tmp_lims)))) {
          dim(tmp_lims) <- c(var = 1, dim(tmp_lims))
        }
        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) 
      }
    }
  }

  # Compute Probabilities
  for (ps in 1:length(categories)) {
    hcst_probs_ev[[ps]] <- GetProbs(ano_hcst_ev_res, time_dim = 'syear',
                              prob_thresholds = NULL,
                              bin_dim_abs = 'bin',
                              indices_for_quantiles = NULL,
                              memb_dim = 'ensemble',
                              abs_thresholds = lims_ano_hcst_tr_res[[ps]],
                              ncores = ncores)
    obs_probs_ev[[ps]] <- GetProbs(ano_obs_ev_res, time_dim = 'syear',
                                   prob_thresholds = NULL,
                                   bin_dim_abs = 'bin',
                                   indices_for_quantiles = NULL,
                                   memb_dim = 'ensemble',
                                   abs_thresholds = lims_ano_obs_tr_res[[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 <- ano_hcst_ev_res 
  ano_obs <- data$obs
  ano_obs$data <- ano_obs_ev_res

  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(list(Subset(hcst_probs_ev[[ps]],
                             along = 'bin', indices = perc, drop = 'all')),
                           probs_hcst)
      probs_obs <- append(list(Subset(obs_probs_ev[[ps]],
                              along = 'bin', indices = perc, drop = 'all')),
                          probs_obs)
      if (!is.null(data$fcst)) {
        probs_fcst <- append(list(Subset(fcst_probs[[ps]],
                              along = 'bin', indices = perc, drop = 'all')),
                             probs_fcst)
      }
      all_names <- c(all_names, name_elem)
    }
  }
  names(probs_hcst) <- all_names
  if (!('var' %in% names(dim(probs_hcst[[1]])))) {
    probs_hcst <- lapply(probs_hcst, function(x) {
                     dim(x) <- c(var = 1, dim(x))
                     return(x)})
  }
  names(probs_obs) <- all_names
  if (!('var' %in% names(dim(probs_obs[[1]])))) {
    probs_obs <- lapply(probs_obs, function(x) {
                     dim(x) <- c(var = 1, dim(x))
                     return(x)})
  }

  if (!is.null(data$fcst)) {
    names(probs_fcst) <- all_names
    if (!('var' %in% names(dim(probs_fcst[[1]])))) {
      probs_fcst <- lapply(probs_fcst, function(x) {
                     dim(x) <- c(var = 1, dim(x))
                     return(x)})
    }
    if (!('syear' %in% names(dim(probs_fcst[[1]])))) {
      probs_fcst <- lapply(probs_fcst, function(x) {
                     dim(x) <- c(syear = 1, dim(x))
                     return(x)})
    }
  }
  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")
    # TODO Forecast
#    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 = lims_ano_hcst_tr_res,
              #                obs_tr = lims_ano_obs_tr_res),
              probs = list(hcst_ev = hcst_probs_ev,
                           obs_ev = obs_probs_ev,
			   fcst = fcst_probs),
              ref_obs_tr = ano_obs_tr_res))
}


## 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

