# Full-cross-val workflow
## This code should be valid for individual months and temporal averages
source("modules/Crossval/R/tmp/GetProbs.R")
source("modules/Crossval/R/tmp/CST_MergeDims.R") 
source("modules/Saving/Saving.R")

Crossval_anomalies <- function(recipe, data) {
  horizon <- tolower(recipe$Analysis$Horizon)
  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 (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 = 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) {
                                  as.array(quantile(as.vector(x), ps,
                                                    na.rm = TRUE))})},
                              prob_lims = list(unlist(categories)),
                              ncores = ncores)$output1
    names(dim(lims_ano_hcst_tr))[1] <- 'bin'
    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 = TRUE))})},
                             prob_lims = list(unlist(categories)),
                             ncores = ncores)$output1
    names(dim(lims_ano_obs_tr))[1] <- 'bin'
    if (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$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$train.dexes)))
    }
    # check order before storing:
    # Add back 'sday' dim if needed or select central 'sday'
    if ('sday' %in% names(dim(ano_hcst_ev))) {
      if (dim(ano_hcst_ev)['sday'] != 1) {
         # To fail because it should not appear
      }
    } else {
      ano_hcst_ev <- InsertDim(ano_hcst_ev, 'sday', len = 1, pos = 1) 
    }
    # remove syear dim from ev objects
    if ('syear' %in% names(dim(ano_hcst_ev))) {
      if (dim(ano_hcst_ev)['syear'] == 1) {
        ano_hcst_ev <- Subset(ano_hcst_ev, along = 'syear', indices = 1,
                              drop = 'selected')
      } 
    }
    # ----
    if ('sday' %in% names(dim(ano_obs_ev))) {
      if (dim(ano_obs_ev)['sday'] > 1) {
        # 'sday' dim to select the central day
        ano_obs_ev <- Subset(ano_obs_ev, along = 'sday',
                             indices = central_day)
      }
    } else {
      ano_obs_ev <- InsertDim(ano_obs_ev, 'sday', len = 1, pos = 1)
    }
    # ----
    if ('sday' %in% names(dim(ano_obs_tr))) {
      if (dim(ano_obs_tr)['sday'] > 1) {
        # 'sday' dim to select the central day
        ano_obs_tr <- Subset(ano_obs_tr, along = 'sday',
                             indices = central_day)
      }
    } else {
      ano_obs_tr <- InsertDim(ano_obs_tr, 'sday', len = 1, pos = 1)
    }
    # rename dimension to be ensemble
    if ('ensemble' %in% names(dim(ano_obs_tr))) {
      if (dim(ano_obs_tr)['ensemble'] == 1) {
        ano_obs_tr <- Subset(ano_obs_tr, along = 'ensemble', indices = 1,
                         drop = 'selected')
      } else {
        stop("why it is longer than 1?")
      }
    }
    if ('syear' %in% names(dim(ano_obs_tr))) { # should be true
      names(dim(ano_obs_tr))[which(names(dim(ano_obs_tr)) == 'syear')] <- 'ensemble'
    }
    # ----
    if ('sday' %in% names(dim(lims_ano_obs_tr))) {
      if (dim(lims_ano_obs_tr)['sday'] > 1) {
        # 'sday' dim to select the central day
        lims_ano_obs_tr <- Subset(lims_ano-obs_tr, along = 'sday',
                              indices = central_day)
      }
    } else {
      lims_ano_obs_tr <- InsertDim(lims_ano_obs_tr, 'sday', len = 1, pos = 1)
    }
    # lims cannot have ensemble dim:
    if ('ensemble' %in% names(dim(lims_ano_obs_tr))) {
      if (dim(lims_ano_obs_tr)['ensemble'] == 1) {
        lims_ano_obs_tr <- Subset(lims_ano_obs_tr, along = 'ensemble', 
                              indices = 1, drop = 'selected')
      } else {
        # error:
        stop("Why it has memb_dim")
      }
    }
    # ----
    if ('sday' %in% names(dim(lims_ano_hcst_tr))) {
      if (dim(lims_ano_hcst_tr)['sday'] > 1) {
        # 'sday' dim to select the central day
        lims_ano_hcst_tr <- Subset(lims_ano_hcst_tr, along = 'sday',
                              indices = central_day)
      }
    } else {
      lims_ano_hcst_tr <- InsertDim(lims_ano_hcst_tr, 'sday', len = 1, pos = 1)
    }
    # lims cannot have ensemble dim:
    if ('ensemble' %in% names(dim(lims_ano_hcst_tr))) {
      if (dim(lims_ano_hcst_tr)['ensemble'] == 1) {
        lims_ano_hcst_tr <- Subset(lims_ano_hcst_tr, along = 'ensemble',
                              indices = 1, drop = 'selected')
      } else {
        # error:
        stop("Why it has memb_dim")
      }
    }   
    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))
  }

  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))),
               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()
  

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

  # Forecast anomalies:
  if (!is.null(data$fcst)) {
    if (horizon %in% c('subseasonal')) {
      # merge sample dimensions and select central week
      hcst <- MergeDims(data$hcst$data, merge_dims = c('sday', 'syear'),
                          rename_dim = 'syear', na.rm = FALSE)
      hcst <- Subset(hcst, along = 'sweek',
                     indices = (dim(data$hcst$data)['sweek'] + 1) / 2)
      obs <- MergeDims(data$obs$data, merge_dims = c('sday', 'syear'),
                          rename_dim = 'syear', na.rm = FALSE)
      obs <- Subset(obs, along = 'sweek',
                    indices = (dim(data$obs$data)['sweek'] + 1) / 2)
      fcst <- Subset(data$fcst$data, along = 'sday', indices = 1)
    } else {
      # we need to keep original objects (to later compute bias)
      hcst <- data$hcst$data
      obs <- data$obs$data
      fcst <- data$fcst$data
    }   
    clim_hcst <- Apply(hcst,
                       target_dims = c('syear', 'ensemble'),
                       mean,
                       na.rm = na.rm,
                       ncores = ncores)$output1
    data$fcst$data <- Ano(data = fcst, clim = clim_hcst)
    hcst_ano <- Ano(data = hcst, 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) {
                                as.array(quantile(as.vector(x), ps,
                                                  na.rm = TRUE))})},
                       output_dims = lapply(categories, function(x) {'bin'}),
                       prob_lims = categories,
                       ncores = ncores)
    clim_obs <- Apply(obs,
                      target_dims = c('syear', 'ensemble'),
                      mean,
                      na.rm = na.rm,
                      ncores = ncores)$output1
    obs_ano <- Ano(data = obs, 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 = TRUE))})},
                  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 (length(probs_fcst) > 0) {
      save_probabilities(recipe = recipe, probs = probs_fcst,
                         data_cube = data$fcst, agg = agg,
                         type = "fcst")
    }
  }
  # Subset dimension 'sday' for subseasonal for mean bias:
  if (horizon == 'subseasonal') {
    central_day <- (dim(data$hcst$data)['sday'] + 1)/2
    data$hcst <- CST_Subset(data$hcst, along = 'sday', indices = central_day)
    data$obs <- CST_Subset(data$obs, along = 'sday', indices = central_day)
  }
  return(list(hcst = ano_hcst, obs = ano_obs, fcst = data$fcst,
              hcst.full_val = data$hcst, obs.full_val = data$obs,
              cat_lims = list(obs = res$lims_ano_obs_tr),
              probs = list(hcst = hcst_probs_ev,
                           obs  = obs_probs_ev,
			   fcst = fcst_probs),
              ref_obs = 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

