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

week <- format(seq(as.Date("20240104", "%Y%m%d"), as.Date("20241216", "%Y%m%d"), 7), "%Y%m%d")
weeknum <- 1:length(week)
weeknum[lubridate::year(as.Date(week, "%Y%m%d")) == 2024] <- weeknum[lubridate::year(as.Date(week, "%Y%m%d")) == 2024] - 52
weeknum <- sprintf("%02d", weeknum)


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

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(Granollers)


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_probs <- paste0(path, var, "/$sdates$/", var, "_", region, "_prob_$sdates$.nc")

    probs <- Start(dat = path_probs,
                   var = c('prob_bn', 'prob_n', 'prob_an'),
                   sdates = week,
                   latitude = values(loc[[1]]),
                   longitude = values(loc[[2]]),
                   time = 'all',
                   return_vars = list(latitude = NULL, longitude = NULL, time = 'sdates'),
                   retrieve = T)          

    as.vector(attributes(probs)$Variables$common$latitude)
    as.vector(attributes(probs)$Variables$common$longitude)
    attributes(probs)$Variables$common$time
    dim(attributes(probs)$Variables$common$time)

    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)
    attributes(lims)$Variables$common$time
    dim(attributes(lims)$Variables$common$time)


    path_skill <- paste0(path, var, "/$sdates$/", var, "_", region,
                         "_skill_week$weeknum$.nc")
    dim(week) <- c(weeks = length(week), sdates = 1)
    rpss <- Start(dat = path_skill,
                  var = c('rpss'),
                  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(rpss)$Variables$common$latitude)
    as.vector(attributes(rpss)$Variables$common$longitude)
    attributes(rpss)$Variables$common$time
    dim(attributes(rpss)$Variables$common$time)


    path_obs <- "/esarchive/recon/ecmwf/era5land/weekly_mean/$var$_f1h/$var$_$sdates$.nc"
    week_larger <- format(seq(as.Date("20240520", "%Y%m%d"), as.Date("20240624", "%Y%m%d"), 7), "%Y%m%d")

    ftimes <- sapply(1:(length(week_larger)-3), FUN = function(x) {
          as.character(format(seq(as.Date(week_larger[x], "%Y%m%d"),
                                  as.Date(week_larger[x+3], "%Y%m%d"),
                                  by = "week"), "%Y%m%d"))})
    dim(ftimes) <- c(time = 4, sdates = 3)

    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 = ftimes,
                 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_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(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))
         
    } #ftime 
  } #var
} #reg
#saveRDS(list(skill = rpss[1,1,,1,1,1,1], obs_probs = obs_probs[,1,1,,i,1,1,1,1,1],
#             fcst_probs = probs[1,,,1,1,i], title = toptitle),
#        file = "test_data_plot_det.RDS")
lapply(resultados, function(x) {print(x[[2]])})

legend <- "/esarchive/scratch/nperez/git/Flor/vitigeoss/2nd_asses/legend_precipitation.png"
files <- list.files(output_path, pattern = 'acpre')
setwd(output_path)

dim(files) <- c(file = length(files))
  Apply(list(files), target_dims = NULL, function(x) {
        system(paste("composite -gravity northeast -geometry +10+10",
                     legend, x, x))})

legend <- "/esarchive/scratch/nperez/git/Flor/vitigeoss/2nd_asses/legend_temperature.png"
files <- list.files(output_path, pattern = 't2')
dim(files) <- c(file = length(files))
  Apply(list(files), target_dims = NULL, function(x) {
        system(paste("composite -gravity northeast -geometry +10+10",
                     legend, x, x))})
files <- list.files(output_path, pattern = 'rswin')
dim(files) <- c(file = length(files))
  Apply(list(files), target_dims = NULL, function(x) {
        system(paste("composite -gravity northeast -geometry +10+10",
                     legend, x, x))})




