rm(list=ls())
library(ncdf4)
library(abind)
library(fields)
library(RColorBrewer)
library(s2dv)
library(ggplot2)
library(startR)
library(multiApply)
output_path <- "/esarchive/scratch/ptrascas/sunset/"

week <- format(seq(as.Date("20240104", "%Y%m%d"),
               as.Date("20241226", "%Y%m%d"), 7), "%Y%m%d")
weeknum <- 1:length(week)
weeknum <- sprintf("%02d", which(week == sdates))

vars <- c("t2", "t02min", "t02max")

resultados <- list()

Barcelona <- c(lat = 41.39, lon = 2.1)
Granollers <- c(lat = 41.61, lon = 2.3)
Viladecans <- c(lat = 41.32, lon = 2.0)
locations <- list(Barcelona, Granollers, Viladecans)

path <- "/esarchive/exp/VITIGEOSS/subseasonal/"
region <- "cat"

for (loc in locations) {
  if (loc == Barcelona) {
    loc_name <- "Barcelona"
  } else if (loc == Granollers) {
    loc_name <- "Granollers"
  } else if (loc == Viladecans) {
    loc_name <- "Viladecans"
  }

  for (var in vars) {

    path_lims <- paste0(path, var, "/$sdates$/", var, "_", region,
                        "_percentiles_week$weeknum$.nc")
    dim(week) <- c(weeks = length(week), sdates = 1)
    lims <- Start(dat = path_lims,
                  var = c('p33', 'p66'),
                  sdates = week,
                  weeknum = 'all',
                  weeknum_depends = 'sdates',
                  latitude = values(loc[[1]]),
                  longitude = values(loc[[2]]),
                  time = 'all',
                  return_vars = list(latitude = NULL, longitude = NULL, time = 'sdates'),
                  retrieve = T)        
    as.vector(attributes(lims)$Variables$common$latitude)
    as.vector(attributes(lims)$Variables$common$longitude)
    times <- attributes(lims)$Variables$common$time
    dim(attributes(lims)$Variables$common$time)

    path_obs <- "/esarchive/recon/ecmwf/era5land/weekly_mean/$var$_f1h/$var$_$sdates$.nc"
    mondays <- c("0520", "0517", "0603", "0610", "0617", "0624", "0701","0708","0715","0722")
    mondays <- unlist(lapply(1999:2016, function(x) {paste0(x, mondays)}))
    dim(mondays) <- c(weeks = 10, years = length(1999:2016))

    if (var == "t2") {
      variable <- 'tas'
    } else if (var == "t02min") {
      variable <- 'tasmin'
      path_obs <- "/esarchive/recon/ecmwf/era5land/weekly_mean/$var$_f24h/$var$_$sdates$.nc"
    } else if (var == "t02max") {
      variable <- 'tasmax'
      path_obs <- "/esarchive/recon/ecmwf/era5land/weekly_mean/$var$_f24h/$var$_$sdates$.nc"
    } else if (var == "acprec") {
      variable <- 'prlr'
    } else {
      variable <- 'rsds'
    } 
    obs <- Start(dat = path_obs,
                 var = variable,
                 sdates = mondays,
                 split_multiselected_dims = TRUE,
                 lat = values(loc[[1]]),
                 lon = values(loc[[2]]),
                 return_vars = list(lat = NULL, lon = NULL),
                 retrieve = T)      

  as.vector(attributes(obs)$Variables$common$lat)
    as.vector(attributes(obs)$Variables$common$lon)
    if (variable %in% c('tas', 'tasmin', 'tasmax')) {
      obs <- obs -273.15
    } else if (variable == "prlr") {
      obs <- obs * 3600 * 24 * 1000 *7
    }
obs_lims <- Apply(list(obs), target_dims = 'years',
                  function(x) {quantile(x, c(1/3,2/3))})$output1

lims_formatted <- sprintf("%.2f", obs_lims)
print(paste(loc_name,var,lims_formatted))


obs_probs <- Apply(list(obs, lims), target_dim = list(NULL, 'var'),
                   fun = function(x, y) {
                     if (is.na(x) | is.na(y[1]) | is.na(y[2])) {
                       res <- rep(NA, 3)
                     } else if (x > y[2]) {
                       res <- c(0,0,1)
                     } else if (x < y[1]) {
                       res <- c(1,0,0)
                     } else {
                       res <- c(0,1,0)
                     }
                     return(res)}, output_dims = 'bin')$output1

                     
source("/esarchive/scratch/ptrascas/R/sunset/scripts/probabilistic_forecast.R")

    if (loc == 'Barcelona') {
      place_name <- 'Barcelona'
    } else if (loc == 'Granollers') {
      place_name <- 'Granollers'
    } else if (loc == 'Viladecans') {
      place_name <- 'Viladecans'
    }

    names(dim(obs_probs))[2] <- 'bin'
    for (i in 1:4) {  
      names_x <- as.Date(ftimes[i,], "%Y%m%d")
      colores <- c("Blues", "Greys", "YlOrRd")
      if (var == "t2") {
        variable <- "2-metre air temperature"
      } else if (var == "t02min") {
        variable <- "Minimum 2 metre temperature"
      } else if (var == "t02max") {
        variable <- "Maximum 2 metre temperature"
      } else if (var == "acprec") {
        variable <- "Accumulated precipitation"
        colores <- c("YlOrBr", "Greys", "Greens")
      } else {
        variable <- "Solar radiation"
      }
      toptitle <- paste0(place_name, " - week ", i ,"\n 2024;", variable)
      filename <- paste0(output_path, names(regions)[reg], "_ftime", i, "_", var, ".png")
      most <- plot_probabilistic_forecast(obs_probs[,1,1,i,,1,1,1,1,1], fcst_probs = probs[1,,,1,1,i],
                                  skill = rpss[1,1,,1,1,1,i],
                                  thresholds = NULL, units = NULL, title = toptitle, title_line = -1,
                                  cols = colores ,#c("Blues", "YlOrBr", "Reds"),
                                  category_names = c('Above normal','Normal','Below normal'),
                                  xlabs = names_x, style = 'most-likely-skill', margins = c(0,4,0,0),
                                  cat_dim = 'bin', time_dim = 'sdates', auxiliar_lines = TRUE,
                                  width = 2000, height = 1000, res = 200, cut_figure = TRUE, 
                                  fileout = filename)
     hits <-  Apply(most, target_dims = list('bin', NULL), function(x, y) {
            if (y %in% which(x == 0.9)) { 
              res <- 1
            } else if (y %in% which(x == 0.6)) {
              res <- 0.5
            } else {
              res <- 0
            }
            return(res)}
            )$output1
      hits_skilled <- sum(hits == 1, na.rm = T)
      hits <- sum(hits >= 0.5, na.rm = T)
      observations <- sum(!is.na(most$observed))
      results <- list(filename, hits_skilled = hits_skilled, hits = hits, observations = observations)
print(results)
      resultados <- append(resultados, list(results))

	  }}}














