# 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/NAO.R")  
source("modules/Crossval/R/tmp/ProjectField.R")
source("modules/Crossval/R/tmp/EOF.R")
source("modules/Crossval/R/tmp/Utils.R")

Crossval_NAO <- function(recipe, data) {
  obsproj <- recipe$Analysis$Workflow$Indices$NAO$obsproj
  if (is.null(obsproj)) {
    obsproj <- TRUE
  }
  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
  nlats <- dim(data$hcst$data)['latitude']
  nlons <- dim(data$hcst$data)['longitude']
  agg = '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 = ?

  ## output objects
  nao_hcst_ev_res <- NULL
  nao_obs_ev_res <- NULL
  nao_obs_tr_res <- NULL
  # as long as probs requested in recipe:
  lims_nao_hcst_tr_res <- lapply(categories, function(X) {NULL})
  lims_nao_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)
    obs_ev <- Subset(data$obs$data, along = 'syear',
  		     indices = cross[[t]]$eval.dexes)
    # 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 NAO:
    nao <- NAO(exp = ano_hcst_tr, obs = ano_obs_tr, exp_cor = ano_hcst_ev,
               ftime_avg = NULL, time_dim = 'syear',
               memb_dim = 'ensemble',
               space_dim = c('latitude', 'longitude'),
               ftime_dim = 'time', obsproj = obsproj,
               lat = data$obs$coords$latitude,
               lon = data$obs$coords$longitude, 
               ncores = recipe$Analysis$ncores)

    nao_obs_ev <- NAO(exp = ano_hcst_tr, obs = ano_obs_tr, exp_cor = ano_obs_ev,
                      ftime_avg = NULL, time_dim = 'syear',
                      memb_dim = 'ensemble',
                      space_dim = c('latitude', 'longitude'),
                      ftime_dim = 'time', obsproj = obsproj,
                      lat = data$obs$coords$latitude,
                      lon = data$obs$coords$longitude,
                      ncores = recipe$Analysis$ncores)$exp_cor
    #Standarisation:
    # Need the nao_hcst (for the train.dexes) to standarize the eval.dexes?
    nao_hcst_ev <- Apply(list(nao$exp, nao$exp_cor),
                         target_dims = c('syear', 'ensemble'),
                         fun = function(x, y) {
                           sd <- sqrt(var(as.vector(x), na.rm = TRUE))
                           means <- mean(as.vector(x), na.rm = TRUE)
                           res <- apply(y, c(1,2), function(z) {(z-means)/sd})},
                         ncores = recipe$Analysis$ncores)$output1
    nao_obs_ev <- Apply(list(nao$obs, nao_obs_ev),
                        target_dims = c('syear','ensemble'),
                        fun = function(x, y) {
                          sd <- sqrt(var(as.vector(x), na.rm = TRUE))
                          means <- mean(as.vector(x), na.rm = TRUE)
                          res <- apply(y, c(1,2), function(z) {(z-means)/sd})},
                        ncores = recipe$Analysis$ncores)$output1
    nao_obs_tr <- Apply(list(nao$obs), target_dims = 'syear',
                        fun = function(x) {
                          sd <- sqrt(var(as.vector(x), na.rm = TRUE))
                          means <- mean(as.vector(x), na.rm = TRUE)
                          res <- apply(x, 1, function(z) {(z-means)/sd})},
                        ncores = recipe$Analysis$ncores,
                        output_dims = 'syear')$output1
    nao_hcst_tr <- Apply(list(nao$exp), target_dims = c('syear', 'ensemble'),
                         fun = function(x) {
                           sd  <- sqrt(var(as.vector(x), na.rm = TRUE))
                           means <- mean(as.vector(x), na.rm = TRUE)
                           res <- apply(x, c(1,2), function(z) {(z-means)/sd})},
                          ncores = recipe$Analysis$ncores)$output1
 
    # compute category limits
    lims_nao_hcst_tr <- Apply(nao_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) {'cat'}), 
                              prob_lims = categories,
                              ncores = ncores)
    lims_nao_obs_tr <- Apply(nao_obs_tr, target_dims = c('syear'),
                             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){'cat'}),
                             prob_lims = categories,
                             ncores = ncores)
    #store results
    nao_hcst_ev_res <- abind(nao_hcst_ev_res, nao_hcst_ev,
                             along = length(dim(nao_hcst_ev)) + 1)
    nao_obs_ev_res <- abind(nao_obs_ev_res, nao_obs_ev,
                            along = length(dim(nao_obs_ev)) + 1)
    nao_obs_tr_res <- abind(nao_obs_tr_res, nao_obs_tr,
                            along = length(dim(nao_obs_tr)) + 1)
    for(ps in 1:length(categories)) {
      lims_nao_hcst_tr_res[[ps]] <- abind(lims_nao_hcst_tr_res[[ps]],
                                          lims_nao_hcst_tr[[ps]],
                                along = length(dim(lims_nao_hcst_tr[[ps]])) + 1)
      lims_nao_obs_tr_res[[ps]] <- abind(lims_nao_obs_tr_res[[ps]],
                                         lims_nao_obs_tr[[ps]],
                                 along = length(dim(lims_nao_obs_tr[[ps]])) + 1)
    }
  }
  info(recipe$Run$logger,
       "#### Anomalies Cross-validation loop ended #####")
  gc()
  # Add dim names:
  names(dim(nao_hcst_ev_res)) <- c('unneeded', 'ensemble', 'dat', 'var',
                                   'sday', 'sweek', 'time', 'syear')
  names(dim(nao_obs_ev_res)) <- c('unneeded', 'ensemble', 'dat', 'var',
                                   'sday', 'sweek', 'time', 'syear')
  names(dim(nao_obs_tr_res)) <- c('ensemble', 'unneeded', 'dat', 'var',
                                  'sday', 'sweek', 'time', 'syear')
  # To make crps_clim to work the reference forecast need to have same dims as obs:
  nao_hcst_ev_res <- Subset(nao_hcst_ev_res, along = 'unneeded',
                           indices = 1, drop = 'selected')
  nao_obs_ev_res <- Subset(nao_obs_ev_res, along = 'unneeded',
                           indices = 1, drop = 'selected')
  nao_obs_tr_res <- Subset(nao_obs_tr_res, along = 'unneeded',
                           indices = 1, drop = 'selected')
  for(ps in 1:length(categories)) {
    names(dim(lims_nao_hcst_tr_res[[ps]])) <- c('cat', 'dat', 'var', 'sday',
                                                'sweek', 'time', 'syear')
    names(dim(lims_nao_obs_tr_res[[ps]])) <- c('cat', 'unneeded', 'dat', 'var',
                                               'sday', 'sweek', 'time', 'syear')
    lims_nao_obs_tr_res[[ps]] <- Subset(lims_nao_obs_tr_res[[ps]],
                                        along = 'unneeded', indices = 1, drop = 'selected')
  }
  # Forecast anomalies:
  nao_fcst <- NULL
  if (!is.null(data$fcst)) {
    clim_hcst <- Apply(data$hcst$data,
                       target_dims = c('syear', 'ensemble'),
                       mean,
                       na.rm = na.rm,
                       ncores = ncores)$output1
    clim_obs <- Apply(data$obs$data,
                       target_dims = c('syear', 'ensemble'),
                       mean,
                       na.rm = na.rm,
                       ncores = ncores)$output1
    ano_fcst <- Ano(data = data$fcst$data, clim = clim_hcst)
    ano_hcst <- Ano(data = data$hcst$data, clim = clim_hcst)
    ano_obs <- Ano(data= data$obs$data, clim = clim_obs)
    nao_fcst <- NAO(exp = ano_hcst, obs = ano_obs, exp_cor = ano_fcst,
                    ftime_avg = NULL, time_dim = 'syear',
                    memb_dim = 'ensemble',
                    space_dim = c('latitude', 'longitude'),
                    ftime_dim = 'time', obsproj = obsproj,
                    lat = data$obs$coords$latitude,
                    lon = data$obs$coords$longitude,
                    ncores = recipe$Analysis$ncores)
    # Terciles limits using the whole hindcast period:
    lims_fcst <- Apply(nao_fcst$exp, 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) {'cat'}),
                       prob_lims = categories,
                       ncores = ncores)
  }

  # Compute Probabilities
  for (ps in 1:length(categories)) {
    hcst_probs_ev[[ps]] <- GetProbs(nao_hcst_ev_res, time_dim = 'syear',
                              prob_thresholds = NULL,
                              bin_dim_abs = 'cat',
                              indices_for_quantiles = NULL,
                              memb_dim = 'ensemble',
                              abs_thresholds = lims_nao_hcst_tr_res[[ps]],
                              ncores = ncores)
    obs_probs_ev[[ps]] <- GetProbs(nao_obs_ev_res, time_dim = 'syear',
                                   prob_thresholds = NULL,
                                   bin_dim_abs = 'cat',
                                   indices_for_quantiles = NULL,
                                   memb_dim = 'ensemble',
                                   abs_thresholds = lims_nao_obs_tr_res[[ps]],
                                   ncores = ncores)
    if (!is.null(data$fcst)) {
      fcst_probs[[ps]] <- GetProbs(nao_fcst$exp_cor, time_dim = 'syear',
                                   prob_thresholds = NULL,
                                   bin_dim_abs = 'cat',
                                   indices_for_quantiles = NULL,
                                   memb_dim = 'ensemble',
                                   abs_thresholds = lims_fcst[[ps]],
                                   ncores = ncores)
    }
  }
  # Convert to s2dv_cubes
  nao_hcst <- data$hcst
  nao_hcst$data <- nao_hcst_ev_res 
  nao_obs <- data$obs
  nao_obs$data <- nao_obs_ev_res

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

  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 = 'cat', indices = perc, drop = 'all')),
                           probs_hcst)
      probs_obs <- append(list(Subset(obs_probs_ev[[ps]],
                              along = 'cat', indices = perc, drop = 'all')),
                          probs_obs)
      if (!is.null(data$fcst)) {
        probs_fcst <- append(list(Subset(fcst_probs[[ps]],
                              along = 'cat', 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")
    }
  }
  return(list(hcst = nao_hcst, obs = nao_obs, fcst = nao_fcst,
              # Mean Bias for the NAO doesn't make sense, so
              # not returning fullvals
              #hcst.full_val = data$hcst, obs.full_val = data$obs,
              cat_lims = list(hcst_tr = lims_nao_hcst_tr_res,
                              obs_tr = lims_nao_obs_tr_res),
              probs = list(hcst_ev = hcst_probs_ev,
                           obs_ev = obs_probs_ev),
              ref_obs_tr = nao_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

