# Full-cross-val workflow
## This code should be valid for individual months and temporal averages

## data dimensions 
sdate_dim <- dim(data$hcst$data)['syear']
nmemb <- dim(data$hcst$data)['ensemble']
nftime <- dim(data$hcst$data)['time']

cross <- CSTools:::.make.eval.train.dexes('leave-one-out', sdate_dim, NULL)
outdim <- length(cross) # leave-one-out should be equal to sdate_dim

# What we need to return?

if ('latitude' %in% names(dim(data$hcst$data))) {
  nlats <- dim(data$hcst$data)['latitude']
  nlons <- dim(data$hcst$data)['longitude']
} else if  ('region' %in% names(dim(data$hcst$data))) {
  nregions <- dim(data$hcst$data)['region']
}
ano_hcst_ev_res <- NULL
ano_obs_ev_res <- NULL
ano_obs_tr_res <- NULL

cal_hcst_ev_res <- NULL
cal_hcst_tr_res <- NULL
obs_ev_res <- NULL
obs_tr_res <- NULL

lims_ano_hcst_tr_res <- NULL
lims_ano_obs_tr_res <- NULL

lims_cal_hcst_tr_res <- NULL
lims_cal_obs_tr_res <- NULL

for (t in 1:outdim) {
  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)
  obs_ev <- Subset(data$obs$data, along = 'syear',
		   indices = cross[[t]]$eval.dexes)
  if (anomalies_comp) {
    # 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)
    #rm("clim_obs_tr", "clim_hcst_tr", "obs_tr", "hcst_tr", "obs_ev", "hcst_ev")
    #Category limits
    lims_ano_hcst_tr <- Apply(ano_hcst_tr, target_dims = c('syear', 'ensemble'),
                            fun = function(x) {
                            quantile(as.vector(x), c(1/3, 2/3), na.rm = TRUE)},
                            output_dims = 'probs',
                            ncores = recipe$Analysis$ncores)$output1
    lims_ano_obs_tr <- Apply(ano_obs_tr, target_dims = c('syear'),#, 'ensemble'),
                           fun = function(x) {
                           quantile(as.vector(x), c(1/3, 2/3), na.rm = TRUE)},
                           output_dims = 'probs',
                           ncores = recipe$Analysis$ncores)$output1
    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)
    lims_ano_hcst_tr_res <- abind(lims_ano_hcst_tr_res, lims_ano_hcst_tr,
                                along = length(dim(lims_ano_hcst_tr)) + 1)
    lims_ano_obs_tr_res <- abind(lims_ano_obs_tr_res, lims_ano_obs_tr,
                                along = length(dim(lims_ano_obs_tr)) + 1)
  }
  if (calibraion_comp) {
    cal_hcst_tr <- Calibration(hcst_tr, obs_tr, exp_cor = NULL, 
                        cal.method = "mse_min", eval.method = "in-sample",  
                        multi.model = FALSE, na.fill = TRUE, 
                        na.rm = TRUE, apply_to = NULL, alpha = NULL,
                        memb_dim = 'ensemble', sdate_dim = 'syear',
                        dat_dim = NULL, 
                        ncores = recipe$Analysis$ncores)
    cal_hcst_ev <- Calibration(hcst_tr, obs_tr, exp_cor = hcst_ev,
                        cal.method = "mse_min",
                        eval.method = "hindcast-vs-forecast",    
                        multi.model = FALSE, na.fill = TRUE,
                        na.rm = TRUE, apply_to = NULL, alpha = NULL,
                        memb_dim = 'ensemble', sdate_dim = 'syear',
                        dat_dim = NULL,
                        ncores = recipe$Analysis$ncores)
    #Category limits
    lims_cal_hcst_tr <- Apply(cal_hcst_tr, target_dims = c('syear', 'ensemble'),
                            fun = function(x) {
                            quantile(as.vector(x), c(1/3, 2/3), na.rm = TRUE)},
                            output_dims = 'probs',
                            ncores = recipe$Analysis$ncores)$output1
    lims_cal_obs_tr <- Apply(obs_tr, target_dims = c('syear'),#, 'ensemble'),
                           fun = function(x) {
                           quantile(as.vector(x), c(1/3, 2/3), na.rm = TRUE)},
                           output_dims = 'probs',
                           ncores = recipe$Analysis$ncores)$output1
    cal_hcst_tr <- Subset(cal_hcst_tr, indices = 1, along = 'syear', drop = 'selected')
    cal_hcst_tr_res <- abind(cal_hcst_tr_res, cal_hcst_tr,
                           along = length(dim(cal_hcst_tr)) + 1)
    names(dim(cal_hcst_tr_res)) <- c(names(dim(cal_hcst_tr)), names(sdate_dim))
    cal_hcst_ev <- Subset(cal_hcst_ev, indices = 1, along = 'syear', drop = 'selected')
    cal_hcst_ev_res <- abind(cal_hcst_ev_res, cal_hcst_ev,
                           along = length(dim(cal_hcst_ev)) + 1)
    names(dim(cal_hcst_ev_res)) <- c(names(dim(cal_hcst_ev)), names(sdate_dim))

    obs_tr_res <- abind(obs_tr_res, obs_tr,
                           along = length(dim(obs_tr)) + 1)
    names(dim(obs_tr_res)) <- c(names(dim(obs_tr)), names(sdate_dim))
    lims_cal_hcst_tr_res <- abind(lims_cal_hcst_tr_res, lims_cal_hcst_tr,
                                along = length(dim(lims_cal_hcst_tr)) + 1)
    names(dim(lims_cal_hcst_tr_res)) <- c(names(dim(lims_cal_hcst_tr)), names(sdate_dim))
    lims_cal_obs_tr <- Subset(lims_cal_obs_tr, indices = 1, along = 'ensemble', drop = 'selected')
    lims_cal_obs_tr_res <- abind(lims_cal_obs_tr_res, lims_cal_obs_tr,
                                along = length(dim(lims_cal_obs_tr)) + 1)
    names(dim(lims_cal_obs_tr_res)) <- c(names(dim(lims_cal_obs_tr)), names(sdate_dim))
  }
  gc()
  res <- list(ano_hcst_ev = ano_hcst_ev_res,
              ano_obs_ev = ano_obs_ev_res,
              ano_obs_tr = ano_obs_tr_res, #reference forecast for the CRPSS
              lims_ano_hcst_tr = lims_ano_hcst_tr_res,
              lims_ano_obs_tr = lims_ano_obs_tr_res,
              cal_hcst_tr = cal_hcst_tr_res,
              cal_hcst_ev = cal_hcst_ev_res,
              obs_tr = obs_tr_res,
              lims_cal_hcst_tr = lims_cal_hcst_tr_res,
              lims_cal_obs_tr = lims_cal_obs_tr_res )
}
info(recipe$Run$logger, 
     paste0("Cross-validation loop ended, returning elements:", 
	    paste(names(res), collapse = " ")))




