source("modules/Loading/Loading.R")
source("modules/Saving/Saving.R")
source("modules/Units/Units.R")
source("modules/Visualization/Visualization.R")
source("modules/Skill/Skill.R")
args = commandArgs(trailingOnly = TRUE)
#recipe_file <- args[1]
recipe_file <- "recipe_bigpredidata_fqa.yml"
#recipe <- read_atomic_recipe(recipe_file)
recipe <- prepare_outputs(recipe_file)
# 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))



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")
             cal_hcst_tr <- CSTools::Calibration(exp = hcst_tr, obs = obs_tr,
                                        cal.method = recipe$Analysis$Workflow$Calibration$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 = recipe$Analysis$Workflow$Calibration$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

             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
   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)
} 
# RPS
source("GetProbs.R")
hcst_probs_ev <- GetProbs(res$hcst_ev, time_dim = 'syear', 
                 prob_thresholds = NULL,
                 bin_dim_abs = 'probs', 
		 	      indices_for_quantiles = NULL,
                 memb_dim = 'ensemble', abs_thresholds = res$lims_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("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)
# CRPS
crps <- CRPS(exp  = res$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 información 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$hcst_ev, obs = res$obs_ev, ref = ref,
               memb_dim = 'ensemble',
               time_dim = 'syear', clim.cross.val = FALSE,
               ncores = recipe$Analysis$ncores)


# Corr
source("Corr.R")
enscorr <- Corr(res$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$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("SprErr.R")
enssprerr <- SprErr(exp = res$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$sing))) {
      info(recipe$Run$logger,
       "RPSS NA")

       rpss$sing[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, #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)
      save_metrics(recipe = recipe,
                   skill = skill_metrics,
                   data_cube = data$hcst, agg = 'global',
                   outdir = recipe$Run$output_dir)

recipe$Run$output_dir <- original

source("modules/Visualization/Visualization.R")
if (data$hcst$coords$longitude[1] != 0) {
  skill_metrics <- lapply(skill_metrics, function(x) {
                Subset(x, along = 'longitude', indices = c(182:360, 1:181))
        })
}
     info(recipe$Run$logger,
       paste("lons:", data$hcst$coords$longitude))
     info(recipe$Run$logger,
       paste("lons:", data$obs$coords$longitude))


#data$hcst$coords$longitude <- -179:180

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

source("tools/add_logo.R")
add_logo(recipe, "rsz_rsz_bsc_logo.png")



