
source("modules/Loading/Loading.R")
#source("modules/Units/Units.R")
source("modules/Saving/Saving.R")
source("modules/Visualization/Visualization.R")
args = commandArgs(trailingOnly = TRUE)
recipe_file <- args[1]
#recipe_file <- "recipe_NAO_scorecards.yml"
recipe <- read_atomic_recipe(recipe_file)
#recipe <- prepare_outputs(recipe_file)
# Load datasets
data <- Loading(recipe)
#data <- Units(recipe, data)
# Full-cross-val workflow
sdate_dim <- dim(data$hcst$data)['syear']
cross <- CSTools:::.make.eval.train.dexes('leave-one-out', sdate_dim, NULL)
# Paralelized:
loops <- array(1:length(cross), c(loop = length(cross)))

 source("/esarchive/scratch/nperez/git/s2dv/R/NAO.R")
 source("/esarchive/scratch/nperez/git/s2dv/R/Utils.R")
 source("/esarchive/scratch/nperez/git/s2dv/R/EOF.R")
 source("/esarchive/scratch/nperez/git/s2dv/R/ProjectField.R")


res <- Apply(list(loops), target = NULL,
             fun = function(t) {
             # subset years:
             # training
             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)
             # eval years
             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)
             ano_hcst_tr <- s2dv::Ano(hcst_tr, clim_hcst_tr) 
             ano_hcst_ev <- s2dv::Ano(hcst_ev, clim_hcst_tr)
             ano_obs_ev <- s2dv::Ano(obs_ev, clim_obs_tr)
             # 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',
                        lat = data$obs$attrs$Variable$metadata$lat,
                        lon = data$obs$attrs$Variable$metadata$lon)

             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',
                               lat = data$obs$attrs$Variable$metadata$lat,
                               lon = data$obs$attrs$Variable$metadata$lon)$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 = 1)$output1
             nao_obs_ev <- Apply(list(nao$obs, nao_obs_ev),
                                   target_dims = list('syear', 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 = 1)$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 = 1, 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 = 1)$output1
             #Category limits
             lims_nao_hcst_tr <-  Apply(nao_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')$output1
             lims_nao_obs_tr <-  Apply(nao_obs_tr, target_dims = 'syear',
                   fun = function(x) {quantile(as.vector(x), c(1/3, 2/3), na.rm = TRUE)},
                   output_dims = 'probs')$output1

             return(list(#ano_obs_tr = ano_obs_tr, ano_hcst_tr = ano_hcst_tr,
                         #ano_hcst_ev = ano_hcst_ev, ano_obs_ev = ano_obs_ev,
                         #clim_obs_tr = clim_obs_tr, clim_hcst_tr = clim_hcst_tr,
                         nao_hcst_tr = nao_hcst_tr, nao_hcst_ev = nao_hcst_ev,
                         nao_obs_ev = nao_obs_ev, nao_obs_tr = nao_obs_tr,
                         lims_nao_hcst_tr = lims_nao_hcst_tr, lims_nao_obs_tr = lims_nao_obs_tr))
              }, ncores = recipe$Analysis$ncores)
# RPS
source("/esarchive/scratch/nperez/git2/s2dv/R/GetProbs.R")
nao_hcst_probs_ev <- GetProbs(res$nao_hcst_ev, time_dim = 'loop', 
                 prob_thresholds = NULL,
                 prob_dim = 'probs', indices_for_quantiles = NULL,
                 memb_dim = 'ensemble', abs_thresholds = res$lims_ano_hcst_tr,
                 ncores = recipe$Analysis$ncores)
nao_obs_probs_ev <- GetProbs(res$nao_obs_ev, time_dim = 'loop',
                 prob_thresholds = NULL,
                 prob_dim = 'probs', indices_for_quantiles = NULL,
                 memb_dim = 'ensemble', abs_thresholds = res$lims_ano_obs_tr,
                 ncores = recipe$Analysis$ncores)
rps <- RPS(exp = nao_hcst_probs_ev, obs = nao_obs_probs_ev, memb_dim = NULL,
           cat_dim = 'probs', cross.val = FALSE, time_dim = 'loop',
           ncores = recipe$Analysis$ncores)
# RPSS
rpss <- RPSS(exp = nao_hcst_probs_ev, obs = nao_obs_probs_ev,
                   time_dim = 'loop', memb_dim = NULL,
                   cat_dim = 'probs',
                   # We should use a line like this 
                   #abs_threshold = res$lims_ano_hcst_tr, 
                   #prob_threshold = c(1/3, 2/3),
                   cross.val = FALSE, 
                   ncores = recipe$Analysis$ncores)
# CRPS
crps <- CRPS(exp  = res$nao_hcst_ev, obs = res$nao_obs_ev,
             time_dim = 'loop', memb_dim = 'ensemble',
             ncores = recipe$Analysis$ncores)
# CRPSS
crpss <- CRPSS(exp = res$nao_hcst_ev, obs = res$nao_obs_ev, 
               memb_dim = 'ensemble',
               time_dim = 'loop', clim.cross.val = TRUE,
               ncores = recipe$Analysis$ncores)

# Corr
enscorr <- s2dv::Corr(res$nao_hcst_ev, res$nao_obs_ev, 
                          dat_dim = 'dat', 
                          time_dim = 'loop', 
                          method = 'pearson', 
                          memb_dim = 'ensemble',
                          memb = F,
                          conf = F,
                          pval = F,
                          sign = T,
                          alpha = 0.05,
                          ncores = recipe$Analysis$ncores)

# Mean Bias
#mean_bias <- Bias(res$ano_hcst_ev, res$ano_obs_ev,
mean_bias <- Bias(data$hcst$data, data$obs$data,
                 time_dim = 'syear',
                 memb_dim = 'ensemble',
                 ncores = recipe$Analysis$ncores)
# Spread error ratio
obs_noensdim <- ClimProjDiags::Subset(res$ano_obs_ev, "ensemble", 1,
                                            drop = "selected")
enssprerr <- easyVerification::veriApply(verifun = 'EnsSprErr',
                                         fcst = res$ano_hcst_ev,
                                         obs = obs_noensdim,
                                         tdim = which(names(dim(res$ano_hcst_ev))=='loop'),
                                         ensdim = which(names(dim(res$ano_hcst_ev))=='ensemble'),
                                         na.rm = TRUE,
                                         ncpus = recipe$Analysis$ncores)

skill_metrics <- list(mean_bias = mean_bias, enscorr = enscorr$corr,
                      enscorr_significance = enscorr$sign, enssprerr = enssprerr,
                      #rps = rps,
                      rpss = rpss$rpss, rpss_significance = rpss$sign, #crps = crps,
                      crpss = crpss$crpss, crpss_significance = crpss$sign)
skill_metrics <- lapply(skill_metrics, function(x) {
                      InsertDim(drop(x), len = 1, pos = 1, name = 'var')})
original <- recipe$Run$output_dir
recipe$Run$output_dir <- paste0(original, "/outputs/Skill/")
# Compute save metrics 
source("modules/Saving/Saving.R")
Saving <- Saving(recipe = recipe, data = data, skill = skill_metrics)
recipe$Run$output_dir <- original

source("modules/Visualization/Visualization.R")
#PlotEquiMap(mean_bias[1,1,1,1,1,,c(182:360, 1:181)], 
#            lon = -179:180,
#            lat = data$hcst$attrs$Variable$metadata$latitude, filled.co = F,
#            fileout = "/esarchive/scratch/nperez/test.png")
skill_metrics <- lapply(skill_metrics, function(x) {
              Subset(x, along = 'longitude', indices = c(182:360, 1:181))
      })
data$hcst$coords$longitude <- -179:180
Visualization(recipe, data, skill_metrics, significance = TRUE)
