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

  # Subseasonal: merge sample dimensions and select central week
  if (horizon  %in% c('subseasonal')) {
    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,
                   drop = 'selected')
  } else {
    # we need to keep original objects (to later compute bias)
    hcst <- data$hcst$data
    obs <- data$obs$data
    fcst <- data$fcst$data
  }

  # Limits calculation with full hindcast period
  lims <- list()
  lims <- Apply(obs, 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)

  # Forecast anomalies:
  if (!is.null(data$fcst)) {
    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)

    # Percentile limits for the forecast:
    lims_fcst <- Apply(hcst_cal, 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)
  }

  # 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

  # 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)
    ano_hcst <- CST_Subset(ano_hcst, along = 'sday', indices = central_day)
    ano_obs <- CST_Subset(ano_obs, along = 'sday', indices = central_day)
  }

  data <- list(hcst = ano_hcst, obs = ano_obs, fcst = data$fcst,
               hcst.full_val = data$hcst, obs.full_val = data$obs,
               cat_lims = list(obs = lims),
               probs = list(hcst = hcst_probs_ev,
                            obs  = obs_probs_ev,
                            fcst = fcst_probs),
               ref_obs = res$ano_obs_tr)

  # Save Crossval module outputs
  source("modules/Saving/Saving_crossval.R")
  Saving(recipe, data, module = "Anomalies", agg = agg)

  return(data)
}


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

