source("/esarchive/scratch/ptrascas/R/sunset/modules/Loading/Loading.R")
source("/esarchive/scratch/ptrascas/R/sunset/modules/Saving/Saving.R")
source("/esarchive/scratch/ptrascas/R/sunset/modules/Units/Units.R")
source("/esarchive/scratch/ptrascas/R/sunset/modules/Visualization/Visualization.R")
source("/esarchive/scratch/ptrascas/R/sunset/modules/Skill/Skill.R")
args = commandArgs(trailingOnly = TRUE)
recipe_file <- args[1]
#recipe_file <- "recipes/recipe_bigpredidata_fqa_raw_iberia_tas.yml" # para debug
recipe <- read_atomic_recipe(recipe_file)
#recipe <- prepare_outputs(recipe_file) # para debug
# Load datasets
data <- Loading(recipe)
data <- Units(recipe, data)
data_summary(data$hcst, recipe)
data_summary(data$obs, recipe)

# start from here after doing previous steps with SUNSET recipe:

# Full-cross-val workflow
sdate_dim <- dim(data$hcst$data)['syear']
nmemb <- dim(data$hcst$data)['ensemble']
nftime <- dim(data$hcst$data)['time']
nlats <- dim(data$hcst$data)['latitude']
nlons <- dim(data$hcst$data)['longitude']


info(recipe$Run$logger,
       paste("ftime", nftime))
info(recipe$Run$logger,
       paste(dim(data$obs$data)))
info(recipe$Run$logger,
       paste(as.vector(data$hcst$coords$longitude)))
info(recipe$Run$logger,
       paste(as.vector(data$obs$coords$longitude)))
info(recipe$Run$logger,
       paste(as.vector(data$hcst$coords$latitude)))
info(recipe$Run$logger,
       paste(as.vector(data$obs$coords$latitude)))



cross <- CSTools:::.make.eval.train.dexes('leave-one-out', sdate_dim, NULL)
# Paralelized:
#loops <- array(1:length(cross), c(loop = length(cross)))
#res <- Apply(list(loops), target = NULL,
#             fun = function(t) {
hcst_ev_res <- array(NA, c(nftime, nlats, 
                         nlons, nmemb, 
                         sdate_dim))
cal_hcst_ev_res <- array(NA, c(nftime, nlats,
                         nlons, nmemb,
                         sdate_dim))
obs_ev_res <- array(NA, c(nftime, nlats,
                         nlons, ensemble = 1, sdate_dim))
obs_tr_res <- array(NA, c(sample = sdate_dim - 1, nftime,
                         nlats, nlons, ensemble = 1, sdate_dim))
lims_hcst_tr_res <- array(NA, c(probs = 2, nftime, nlats,
                         nlons, sdate_dim))
lims_obs_tr_res <- array(NA, c(probs = 2, nftime, nlats,
                         nlons, sdate_dim))
lims_cal_hcst_tr_res <- array(NA, c(probs = 2, nftime, nlats,
                         nlons, sdate_dim))

lims_brier10_hcst_tr_res <- array(NA, c(probs = 1, nftime, nlats,
                         nlons, sdate_dim))
lims_brier10_obs_tr_res <- array(NA, c(probs = 1, nftime, nlats,
                         nlons, sdate_dim))
lims_brier10_cal_hcst_tr_res <- array(NA, c(probs = 1, nftime, nlats,
                         nlons, sdate_dim))
lims_brier90_hcst_tr_res <- array(NA, c(probs = 1, nftime, nlats,
                         nlons, sdate_dim))
lims_brier90_obs_tr_res <- array(NA, c(probs = 1, nftime, nlats,
                         nlons, sdate_dim))
lims_brier90_cal_hcst_tr_res <- array(NA, c(probs = 1, nftime, nlats,
                         nlons, sdate_dim))



info(recipe$Run$logger,
       paste(dim(data$hcst$data)))
info(recipe$Run$logger,
       paste(names(dim(data$hcst$data))))

       

#crps_clim_res <- array(NA, c(nftime, nlats, nlons, ensemble = 1, sdate_dim))
for (t in 1:sdate_dim) {
    info(recipe$Run$logger,
       paste("crossval:", t))

             # subset years: Subset works at BSC not at Athos
             # training
             obs_tr <- data$obs$data[1,1,1,1,cross[[t]]$train.dexes,,,,]
             hcst_tr <- data$hcst$data[1,1,1,1,cross[[t]]$train.dexes,,,,]
             # eval years
             hcst_ev <- data$hcst$data[1,1,1,1,cross[[t]]$eval.dexes,,,,]
             obs_ev <- data$obs$data[1,1,1,1,cross[[t]]$eval.dexes,,,,]
             dim(obs_tr) <- c(syear = (as.numeric(sdate_dim) - 1), nftime, nlats, nlons)
	          dim(obs_ev) <- c(syear = 1, nftime, nlats, nlons)
		       dim(hcst_tr) <- c(syear = (as.numeric(sdate_dim) - 1),
		       		              nftime, nlats, nlons, nmemb)
             dim(hcst_ev) <- c(syear = 1, nftime, nlats, nlons, nmemb)
	          
             # rm("obs_tr", "hcst_tr", "obs_ev", "hcst_ev")
             calib.method <- recipe$Analysis$Workflow$Calibration$method
             if (calib.method == 'raw') {
             cal_hcst_tr <- hcst_tr
             cal_hcst_ev <- hcst_ev
             } else {
             cal_hcst_tr <- CSTools::Calibration(exp = hcst_tr, obs = obs_tr,
                                        cal.method = calib.method,
                                        eval.method = 'in-sample', sdate_dim = 'syear',
                                        memb_dim = 'ensemble')
             cal_hcst_ev <- CSTools::Calibration(exp = hcst_tr, obs = obs_tr,
                                        exp_cor = hcst_ev,
                                        cal.method = calib.method,
                                        eval.method = 'in-sample', sdate_dim = 'syear',
                                        memb_dim = 'ensemble')
             }

             #Category limits
             lims_hcst_tr <- Apply(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_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
             lims_cal_hcst_tr <- Apply(cal_hcst_tr, target_dims = c('syear', 'ensemble'),#, 'ensemble'),
                   fun = function(x) {quantile(as.vector(x), c(1/3, 2/3), na.rm = TRUE)},
                   output_dims = 'probs', ncores = recipe$Analysis$ncores)$output1
		   
            # quitamos la categoria "probs" si solo hay una categoria cuando hacemos el brier

             lims_brier10_hcst_tr <- Apply(hcst_tr, target_dims = c('syear', 'ensemble'),
                   fun = function(x) {
		      quantile(as.vector(x), c(1/10), na.rm = TRUE)},
                   ncores = recipe$Analysis$ncores)$output1
             lims_brier10_obs_tr <- Apply(obs_tr, target_dims = c('syear'),#, 'ensemble'),
                   fun = function(x) {quantile(as.vector(x), c(1/10), na.rm = TRUE)},
                   ncores = recipe$Analysis$ncores)$output1
             lims_brier10_cal_hcst_tr <- Apply(cal_hcst_tr, target_dims = c('syear', 'ensemble'),#, 'ensemble'),
                   fun = function(x) {quantile(as.vector(x), c(1/10), na.rm = TRUE)},
                   ncores = recipe$Analysis$ncores)$output1
             lims_brier90_hcst_tr <- Apply(hcst_tr, target_dims = c('syear', 'ensemble'),
                   fun = function(x) {quantile(as.vector(x), c(9/10), na.rm = TRUE)},
                   ncores = recipe$Analysis$ncores)$output1
             lims_brier90_obs_tr <- Apply(obs_tr, target_dims = c('syear'),#, 'ensemble'),
                   fun = function(x) {quantile(as.vector(x), c(9/10), na.rm = TRUE)},
                   ncores = recipe$Analysis$ncores)$output1
             lims_brier90_cal_hcst_tr <- Apply(cal_hcst_tr, target_dims = c('syear', 'ensemble'),#, 'ensemble'),
                   fun = function(x) {quantile(as.vector(x), c(9/10), na.rm = TRUE)},
                   ncores = recipe$Analysis$ncores)$output1




             gc()
   hcst_ev_res[,,,,t] <- hcst_ev
   cal_hcst_ev_res[,,,,t] <- cal_hcst_ev
   obs_ev_res[,,,,t] <- obs_ev
   obs_tr_res[,,,,,t] <- obs_tr
   lims_hcst_tr_res[,,,,t] <- lims_hcst_tr
   lims_cal_hcst_tr_res[,,,,t] <- lims_cal_hcst_tr
   lims_obs_tr_res[,,,,t] <- lims_obs_tr
   lims_brier10_hcst_tr_res[,,,,t] <- lims_brier10_hcst_tr
   lims_brier10_cal_hcst_tr_res[,,,,t] <- lims_brier10_cal_hcst_tr
   lims_brier10_obs_tr_res[,,,,t] <- lims_brier10_obs_tr
   lims_brier90_hcst_tr_res[,,,,t] <- lims_brier90_hcst_tr
   lims_brier90_cal_hcst_tr_res[,,,,t] <- lims_brier90_cal_hcst_tr
   lims_brier90_obs_tr_res[,,,,t] <- lims_brier90_obs_tr



   res <- list(hcst_ev = hcst_ev_res,
               obs_ev = obs_ev_res,
               obs_tr = obs_tr_res, #required as reference forecast for the CRPSS
               cal_hcst_ev = cal_hcst_ev_res,
               lims_hcst_tr = lims_hcst_tr_res,
               lims_obs_tr = lims_obs_tr_res,
               lims_cal_hcst_tr = lims_cal_hcst_tr_res,
	       lims_brier10_hcst_tr = lims_brier10_hcst_tr_res,
	       lims_brier10_obs_tr = lims_brier10_obs_tr_res,
	       lims_brier10_cal_hcst_tr = lims_brier10_cal_hcst_tr_res,
	       lims_brier90_hcst_tr = lims_brier90_hcst_tr_res,
               lims_brier90_obs_tr = lims_brier90_obs_tr_res,
               lims_brier90_cal_hcst_tr = lims_brier90_cal_hcst_tr_res
	       )
        } 
# RPS
source("/esarchive/scratch/ptrascas/R/sunset/scripts/GetProbs.R")
hcst_probs_ev <- GetProbs(res$cal_hcst_ev, time_dim = 'syear', 
                 prob_thresholds = NULL,
                 bin_dim_abs = 'probs', 
		 	      indices_for_quantiles = NULL,
                 memb_dim = 'ensemble', abs_thresholds = res$lims_cal_hcst_tr,
                 ncores = recipe$Analysis$ncores)
obs_probs_ev <- GetProbs(res$obs_ev, time_dim = 'syear',
                 prob_thresholds = NULL,
                 bin_dim_abs = 'probs', 
		 	      indices_for_quantiles = NULL,
                 memb_dim = 'ensemble',
                 abs_thresholds = res$lims_obs_tr,
                 ncores = recipe$Analysis$ncores)
rps <- RPS(exp = hcst_probs_ev, obs = obs_probs_ev, memb_dim = NULL,
           cat_dim = 'probs', cross.val = FALSE, time_dim = 'syear',
           ncores = recipe$Analysis$ncores)
source("/esarchive/scratch/ptrascas/R/sunset/scripts/RPS_clim.R")
rps_clim <- Apply(list(obs_probs_ev),
                  target_dims = c('probs', 'syear'),
                  RPS_clim, bin_dim_abs = 'probs', cross.val = FALSE)$output1
# RPSS
rpss <- RPSS(exp = hcst_probs_ev, obs = obs_probs_ev,
                   time_dim = 'syear', memb_dim = NULL,
                   cat_dim = 'probs',
                   # We should use a line like this 
                   #abs_threshold = res$lims_hcst_tr, 
                   #prob_threshold = c(1/3, 2/3),
                   cross.val = FALSE, 
                   ncores = recipe$Analysis$ncores)

# Brier skill scores
hcst_brier10_probs_ev <- GetProbs(res$cal_hcst_ev, time_dim = 'syear',
                 prob_thresholds = NULL,
                 bin_dim_abs = 'probs',
                 indices_for_quantiles = NULL,
                 memb_dim = 'ensemble', abs_thresholds = res$lims_brier10_cal_hcst_tr,
                 ncores = recipe$Analysis$ncores)
obs_brier10_probs_ev <- GetProbs(res$obs_ev, time_dim = 'syear',
                 prob_thresholds = NULL,
                 bin_dim_abs = 'probs',
                 indices_for_quantiles = NULL,
                 memb_dim = 'ensemble',
                 abs_thresholds = res$lims_brier10_obs_tr,
                 ncores = recipe$Analysis$ncores)
hcst_brier90_probs_ev <- GetProbs(res$cal_hcst_ev, time_dim = 'syear',
                 prob_thresholds = NULL,
                 bin_dim_abs = 'probs',
                 indices_for_quantiles = NULL,
                 memb_dim = 'ensemble', abs_thresholds = res$lims_brier90_cal_hcst_tr,
                 ncores = recipe$Analysis$ncores)
obs_brier90_probs_ev <- GetProbs(res$obs_ev, time_dim = 'syear',
                 prob_thresholds = NULL,
                 bin_dim_abs = 'probs',
                 indices_for_quantiles = NULL,
                 memb_dim = 'ensemble',
                 abs_thresholds = res$lims_brier90_obs_tr,
                 ncores = recipe$Analysis$ncores)


bss10 <- RPSS(exp = hcst_brier10_probs_ev, obs = obs_brier10_probs_ev,
                   time_dim = 'syear', memb_dim = NULL,
                   cat_dim = 'probs', prob_thresholds = c(1/10),
                   ncores = recipe$Analysis$ncores)


bss90 <- RPSS(exp = hcst_brier90_probs_ev, obs = obs_brier90_probs_ev,
                   time_dim = 'syear', memb_dim = NULL,
                   cat_dim = 'probs', prob_thresholds =	c(9/10),
                   ncores = recipe$Analysis$ncores)


## Se puede calcular el Brier Score (no el skill score) usando la formula del RPS
#rps_brier_fcst <- RPS(exp = test_hcst, obs = test_obs, memb_dim = NULL,
#           cat_dim = 'probs', time_dim = 'syear', prob_thresholds = c(1/10),
#           ncores = recipe$Analysis$ncores)

#rps_brier_clim <- Apply(list(test_obs),
#                  target_dims = c('probs', 'syear'),
#                  RPS_clim, bin_dim_abs = 'probs',
#		  prob_thresholds = c(1/10), cross.val = FALSE)$output1



# CRPS
crps <- CRPS(exp  = res$cal_hcst_ev, obs = res$obs_ev,
             time_dim = 'syear', memb_dim = 'ensemble',
             ncores = recipe$Analysis$ncores)
# Este no sé como se calcula????:
# Aquí no se puede porque estaría incluyendo informacion de los otros años
#source("modules/Skill/R/CRPS_clim.R")
# Pero si lo hago con el ano_obs_tr si puedo hacerlo aquí
# el resultado es igual a dentro del bucle.
crps_clim <- CRPS(exp = res$obs_tr, obs = res$obs_ev,
                  time_dim = 'syear', memb_dim = 'sample.syear',
                  ncores = recipe$Analysis$ncores) 


# CRPSS
ref <- res$obs_tr
dim(ref) <- c(ensemble = as.numeric(sdate_dim) -1, 
              nftime, nlats, nlons, sdate_dim)
crpss <- CRPSS(exp = res$cal_hcst_ev, obs = res$obs_ev, ref = ref,
               memb_dim = 'ensemble',
               time_dim = 'syear', clim.cross.val = FALSE,
               ncores = recipe$Analysis$ncores)


# Corr
source("/esarchive/scratch/ptrascas/R/sunset/scripts/Corr.R")
enscorr <- Corr(res$cal_hcst_ev, res$obs_ev, 
                          dat_dim = NULL, 
                          time_dim = 'syear', 
                          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$cal_hcst_ev, res$obs_ev,
mean_bias <- Bias(data$hcst$data, data$obs$data,
                 time_dim = 'syear',
                 memb_dim = 'ensemble',
                 ncores = recipe$Analysis$ncores)

mean_bias_sign <- Apply(list(data$hcst$data, data$obs$data),
  	  		 target_dims = list(c('syear', 'ensemble'), 'syear'),
                         fun = function(x,y) {if (!(any(is.na(x)) || any(is.na(y))))
			 {res <- t.test(x = y, y = apply(x, 1, mean, na.rm = T),
      	          	  alternative = "two.sided")$p.value} else {res <- NA}
			  return(res)},ncores = sdate_dim)$output1
			  
mean_bias_sign <- mean_bias_sign <= 0.05

					     	  
# Spread error ratio
source("/esarchive/scratch/ptrascas/R/sunset/scripts/SprErr.R")
enssprerr <- SprErr(exp = res$cal_hcst_ev, obs = res$obs_ev,
	              memb_dim = 'ensemble', dat_dim = NULL,
                      time_dim = 'syear', pval = TRUE,
		      ncores = recipe$Analysis$ncores)
enssprerr_sign <- enssprerr$p.val
enssprerr_sign <- enssprerr_sign <= 0.05
enssprerr <- enssprerr$ratio


if (any(is.na(rpss$sign))) {
      info(recipe$Run$logger,
       "RPSS NA")

       rpss$sign[is.na(rpss$sign)] <- FALSE
}
skill_metrics <- list(mean_bias = mean_bias,
	      	       mean_bias_significance = !mean_bias_sign,
     		      enscorr = enscorr$corr,
                      enscorr_significance = !enscorr$sign,
		      enssprerr = enssprerr,
                      enssprerr_significance = !enssprerr_sign,
                      rps = rps, rps_clim = rps_clim, crps = crps, crps_clim = crps_clim,
                      rpss = rpss$rpss, rpss_significance = !rpss$sign, 
                      crpss = crpss$crpss, crpss_significance = !crpss$sign,
		      bss10 = bss10$rpss, bss10_significance = !bss10$sign,
		      bss90 = bss90$rpss, bss90_significance = !bss90$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("/esarchive/scratch/ptrascas/R/sunset/modules/Saving/Saving.R")
#Saving <- Saving(recipe = recipe, data = data, skill = skill_metrics)
#save_metrics(recipe = recipe,
#                   skill = skill_metrics,
#                   data_cube = data$hcst, agg = 'global',
#                   outdir = recipe$Run$output_dir)

recipe$Run$output_dir <- original

source("/esarchive/scratch/ptrascas/R/sunset/modules/Visualization/Visualization.R")
     info(recipe$Run$logger,
       paste("lons:", data$hcst$coords$longitude))
     info(recipe$Run$logger,
       paste("lons:", data$obs$coords$longitude))

Visualization(recipe = recipe, data = data, skill_metrics = skill_metrics, significance = TRUE)

#source("/esarchive/scratch/ptrascas/R/sunset/tools/add_logo.R")
#add_logo(recipe, "rsz_rsz_bsc_logo.png")



