rm(list=ls())
library(ncdf4)
library(abind)
library(fields)
library(CSTools)
library(RColorBrewer)
library(s2dv)
library(ggplot2)
library(startR)
library(multiApply)

# Load 1 start date


#args = commandArgs(trailingOnly = TRUE)

sdates <- format(Sys.Date() , "%Y%m%d")

# Nuria's code to load the dates and name the weeks

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


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

# There are some issues picking the location, numbers must be "exact"
Barcelona <- c(lat = 41.39, lon = 2.1)
Granollers <- c(lat = 41.61, lon = 2.29)
Viladecans <- c(lat = 41.31, lon = 2.02)
locations <- list(Barcelona, Granollers)

vars <- c('t2','t02min','t02max') #just for testing

for (loc in locations) {

  if (loc == Barcelona) {
    loc_name <- "Barcelona"
  } else if (loc == Granollers) {
    loc_name <- "Granollers"
  }

    for (var in vars){
    
	# FORECAST
    	path_fcst <- paste0(path, var, "/$sdates$/", var, "_", region, "_$sdates$.nc")
    	fcst <- Start(dat = path_fcst, 
    	          var = var, 
       	     	  sdates = sdates, 
	     	  latitude = values(loc[[1]]), 
             	  longitude = values(loc[[2]]), 
             	  ensemble = 'all', 
             	  time = 'all', 
             	  retrieve = T)

    		  as.vector(attributes(fcst)$Variables$common$latitude)
    		  as.vector(attributes(fcst)$Variables$common$longitude)
    
	# PROBABILITIES
    	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 = 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(probs)$Variables$common$latitude)
    	as.vector(attributes(probs)$Variables$common$longitude)
    	attributes(probs)$Variables$common$time
    	dim(attributes(probs)$Variables$common$time)

    	# TERCILE LIMITS

    	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 = sdates,
                  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)

    	# SKILL

	path_skill <- paste0(path, var, "/$sdates$/", var, "_", region,
                         "_skill_week$weeknum$.nc")

    	rpss <- Start(dat = path_skill,
                  var = c('rpss'),
                  sdates = sdates,
                  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)

	mondays <- format(seq(as.Date(sdates, "%Y%m%d") + 4, as.Date(sdates, "%Y%m%d") + 30, 7), "%d-%m-%Y")
	sundays <- format(seq(as.Date(sdates, "%Y%m%d") + 10, as.Date(sdates, "%Y%m%d") + 40, 7), "%d-%m-%Y")
	subtit <- sapply(1:4, function(x) { paste0("Semana del\n", mondays[x], "\nal ", 
	       	  sundays[x], "\nRPSS = ", round(rpss[1,1,1,1,1,1,x]/100,digits=2))})       

    	# Plot forecast PDF

    	fcst <- data.frame(fcst1 = fcst[1,1,1,1,1,,1], fcst2 = fcst[1,1,1,1,1,,2], fcst3 = fcst[1,1,1,1,1,,3], fcst4 = fcst[1,1,1,1,1,,4])
    	p33 <- c(lims[1,1,1,1,1,1,1],lims[1,1,1,1,1,1,2],lims[1,1,1,1,1,1,3],lims[1,1,1,1,1,1,4])
    	p66 <- c(lims[1,2,1,1,1,1,1],lims[1,2,1,1,1,1,2],lims[1,2,1,1,1,1,3],lims[1,2,1,1,1,1,4])
    	terlims <- array(c(p33,p66), dim = c(4,2)) 

	 # Improve this part to automatise week labels in the subplot titles:

        PlotForecastPDF(fcst,tercile.limits = terlims, var.name = "Temperatura (\u00B0C)",fcst.names = 
    		subtit, title = paste0("Prediccion de temperatura en ",loc_name,"\nEmitida en ", 
		format(as.Date(sdates, "%Y%m%d"), "%d-%m-%Y")), color.set="s2s4e", 
		add.ensmemb = "no", plotfile = paste0("I4C_forecastpdf_",var,"_",loc_name,".png"))
	
     }
   }

print("Hi Palo")

###########################################################################################################################################

# After the operational, add observations for evaluation

# Add observations:

#path_obs <- "/esarchive/recon/ecmwf/era5land/daily_mean/tas_f1h/tas_202308.nc"
#variable <- "tas"
#regions <-  c(lat = 41.39, lon = 2.1)
#obs <- Start(dat = path_obs, var = variable, time = 'all', lat = values(regions[1]), lon = values(regions[2]), return_vars = list(lat = NULL, lon = NULL), retrieve = T)
#obs <- as.s2dv_cube(obs)

#as.vector(attributes(obs)$Variables$common$lat)
#as.vector(attributes(obs)$Variables$common$lon)

# Extract Barcelona and average from day 21 to day 27
#obs <- obs[1,1,21:27,1,1]
#obs <- obs -273.15
#obs <- mean(obs)


#plot <-PlotForecastPDF(fcst,tercile.limits = c(terciles[2],terciles[3]), extreme.limits = c(terciles[1], terciles[4]),
#     var.name = "Temperatura (C)",fcst.names = c(paste0("ftime: days 26-32\n sdate: 27/07/23\nRPSS: ",rpss4),
#     paste0("ftime: days 19-25\n sdate: 03/08/23\nRPSS: ",rpss3),paste0("ftime: days 12-18\n sdate: 10/08/23\nRPSS: ",rpss2),
#     paste0("ftime: days 5-11\n sdate: 17/08/23\nRPSS: ",rpss1)) , obs = c(obs,obs,obs,obs),
#     title = "Prediccion subestacional de la ola de calor de 2023 - Barcelona",color.set="s2s4e")

#ggsave("test1.png", plot, width = 7, height = 5)
