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

if (names(data$hcst$attrs$Variable$metadata) == 'prlr') {
  source("modules/Units/R/deaccumulate.R")
  data$hcst <- deaccumulate(data$hcst, ncores = recipe$Analysis$ncores)
  data_summary(data$hcst, 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) {
ano_hcst_ev_res <- array(NA, c(nftime, nlats, 
                         nlons, nmemb, 
                         sdate_dim))
ano_obs_ev_res <- array(NA, c(nftime, nlats,
                         nlons, ensemble = 1, sdate_dim))
ano_obs_tr_res <- array(NA, c(sample = sdate_dim - 1, nftime,
                         nlats, nlons, ensemble = 1, sdate_dim))
lims_ano_hcst_tr_res <- array(NA, c(probs = 2, nftime, nlats,
                         nlons, sdate_dim))
lims_ano_obs_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,,,,]
	#Subset(data$obs$data, along = 'syear',
        #                          indices = cross[[t]]$train.dexes)
             hcst_tr <- data$hcst$data[1,1,1,1,cross[[t]]$train.dexes,,,,]
	     #Subset(data$hcst$data, along = 'syear',
             #                      indices = cross[[t]]$train.dexes)
             # eval years
             hcst_ev <- data$hcst$data[1,1,1,1,cross[[t]]$eval.dexes,,,,]
             #Subset(data$hcst$data, along = 'syear',
             #                      indices = cross[[t]]$eval.dexes)
             obs_ev <- data$obs$data[1,1,1,1,cross[[t]]$eval.dexes,,,,]
             #Subset(data$obs$data, along = 'syear',
             #                     indices = cross[[t]]$eval.dexes)
       info(recipe$Run$logger,
       paste(names(dim(hcst_ev))))
       info(recipe$Run$logger,
       paste(dim(hcst_ev)))

       info(recipe$Run$logger,
       paste(names(dim(obs_tr))))
       info(recipe$Run$logger,
       paste(dim(obs_tr)))

             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)
	     # compute climatology:
             clim_obs_tr <- MeanDims(obs_tr, 'syear')
             clim_hcst_tr <- MeanDims(hcst_tr, c('syear', 'ensemble'))
             # compute anomalies:
	             info(recipe$Run$logger,
       paste("dims:", dim(clim_hcst_tr)))

             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)

             # CRPS_clim:
             # build the reference forecast usnig the observational trainig sample
             #ref <- array(data = rep(ano_obs_tr[1,1,1,1,,1,1,1,1], each = 1), 
             #             dim = c(sdate = 1, member = sdate_dim))
             #crps_clim <- SpecsVerification::enscrps_cpp(ens = ref, 
             #                  obs = ano_obs_ev[1,1,1,1,,1,1,1,1], R_new = NA)
#             crps_clim <- Apply(list(ano_obs_ev, ano_obs_tr), target_dims = 'syear',
#                                fun = function(x,y) {
#                                  ref <- array(y, c(syear = 1, ensemble = length(y)))
#                                  res <- SpecsVerification::enscrps_cpp(ens = ref,
#                                            obs = x, R_new = NA) #No fair
#                                }, ncores = recipe$Analysis$ncores)$output1
                                  
             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
#mem <- mem_used()
             gc()
   ano_hcst_ev_res[,,,,t] <- ano_hcst_ev
   ano_obs_ev_res[,,,,t] <- ano_obs_ev
   ano_obs_tr_res[,,,,,t] <- ano_obs_tr
   lims_ano_hcst_tr_res[,,,,t] <- lims_ano_hcst_tr
   lims_ano_obs_tr_res[,,,,t] <- lims_ano_obs_tr
#   crps_clim_res[,,,,t] <- crps_clim
   res <- list(ano_hcst_ev = ano_hcst_ev_res,
               ano_obs_ev = ano_obs_ev_res,
               ano_obs_tr = ano_obs_tr_res, #required as reference forecast for the CRPSS
               lims_ano_hcst_tr = lims_ano_hcst_tr_res,
               lims_ano_obs_tr = lims_ano_obs_tr_res)#,
#               crps_clim = crps_clim_res)
} 
#            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,
#                         lims_ano_hcst_tr = lims_ano_hcst_tr,
#                         lims_ano_obs_tr = lims_ano_obs_tr))
#              }, ncores =recipe$Analysis$ncores)
# RPS
source("GetProbs.R")
ano_hcst_probs_ev <- GetProbs(res$ano_hcst_ev, time_dim = 'syear', 
                 prob_thresholds = NULL,
                 bin_dim_abs = 'probs', 
		 indices_for_quantiles = NULL,
                 memb_dim = 'ensemble', abs_thresholds = res$lims_ano_hcst_tr,
                 ncores = recipe$Analysis$ncores)
ano_obs_probs_ev <- GetProbs(res$ano_obs_ev, time_dim = 'syear',
                 prob_thresholds = NULL,
                 bin_dim_abs = 'probs', 
		 indices_for_quantiles = NULL,
                 memb_dim = 'ensemble',
                 abs_thresholds = res$lims_ano_obs_tr,
                 ncores = recipe$Analysis$ncores)
rps <- RPS(exp = ano_hcst_probs_ev, obs = ano_obs_probs_ev, memb_dim = NULL,
           cat_dim = 'probs', cross.val = FALSE, time_dim = 'syear',
           ncores = recipe$Analysis$ncores)
source("modules/Skill/R/RPS_clim.R")
rps_clim <- Apply(list(ano_obs_probs_ev),
                  target_dims = c('probs', 'syear'),
                  RPS_clim, bin_dim_abs = 'probs', cross.val = FALSE)$output1
# RPSS
rpss <- RPSS(exp = ano_hcst_probs_ev, obs = ano_obs_probs_ev,
                   time_dim = 'syear', 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$ano_hcst_ev, obs = res$ano_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$ano_obs_tr, obs = res$ano_obs_ev,
                  time_dim = 'syear', memb_dim = 'sample.syear',
                  ncores = recipe$Analysis$ncores) 


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


# Corr
source("modules/Skill/R/tmp/Corr.R")
enscorr <- Corr(res$ano_hcst_ev, res$ano_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$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)

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(y, apply(x, 1, mean, na.rm = T))$p.value
			          res[res >= 0.5] <- FALSE
			          res[res < 0.5] <- TRUE
				} else {
				  res <- FALSE
				}
				res <- as.logical(res)
			        return(res)},
			ncores = sdate_dim)$output1
		
# 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))=='syear'),
                                         ensdim = which(names(dim(res$ano_hcst_ev))=='ensemble'),
                                         na.rm = FALSE,
                                         ncpus = recipe$Analysis$ncores)
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,
                      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")
#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")
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, "tools/BSC_logo_95.jpg")

