rm(list=ls())

library(CSIndicators)
library(CSTools)
library(zeallot)
library(s2dv)

S5path <- list(name = 'SEAS5',
               path = '/esarchive/exp/ecmwf/system5c3s/$STORE_FREQ$_mean/$VAR_NAME$_s0-24h/$VAR_NAME$_$START_DATE$.nc')

ERA5path <- list(name = 'ERA5',
                 path = '/esarchive/recon/ecmwf/era5/$STORE_FREQ$_mean/$VAR_NAME$_f1h-r1440x721cds/$VAR_NAME$_$YEAR$$MONTH$.nc')

sdates <- paste0(1993:2016, '04', '01')

c(hcst, hcst_ref) %<-% CST_Load(var = 'prlr', 
                                exp = list(S5path), 
                                obs = list(ERA5path), 
                                sdates = sdates, 
                                lonmin = -10, lonmax = 0,
                                latmin = 35, latmax = 40, 
                                storefreq = 'daily', 
                                leadtimemin = 1, leadtimemax = 214, 
                                nmember = 25, output = "lonlat") 
hcst$data <- hcst$data * 3600 * 24 * 1000
attributes(hcst$Variable)$units <- 'mm'
hcst_ref$data <- hcst_ref$data * 3600 * 24 * 1000
attributes(hcst_ref$Variable)$units <- 'mm'


c(fcst, obs) %<-% CST_Load(var = 'prlr',
                           exp = list(S5path),     
                           obs = list(ERA5path),         
                           sdates = '20220401',
                           lonmin = -10, lonmax = 0,
                           latmin = 35, latmax = 40,
                           storefreq = 'daily',
                           leadtimemin = 1, leadtimemax = 214,
                           nmember = 50, output = "lonlat")
fcst$data <- fcst$data * 1000 * 3600 * 24
attributes(fcst$Variable)$units <- 'mm'
obs$data <- obs$data  * 1000 * 3600 * 24
attributes(obs$Variable)$units <- 'mm'


fcst_QM <- CST_QuantileMapping(exp = hcst,
                               obs = hcst_ref,
                               exp_cor = fcst, wet = TRUE,
                               ncores = 6)
hcst_QM <- CST_QuantileMapping(exp = hcst,
                               obs = hcst_ref, wet = TRUE,
                               ncores = 6)

sprR_fcst <- CST_PeriodAccumulation(fcst_QM,
                                    start = list(21,4), 
                                    end = list(21,6), na.rm = FALSE,
                                    time_dim = 'ftime', ncores = 6)
sprR_hcst_ref <- CST_PeriodAccumulation(hcst_ref,
                                    start = list(21,4),
                                    end = list(21,6),
                                    time_dim = 'ftime', ncores = 6)
sprR_obs <- CST_PeriodAccumulation(obs,
                                    start = list(21,4),
                                    end = list(21,6),
                                    time_dim = 'ftime', ncores = 6)
sprR_hcst <- CST_PeriodAccumulation(hcst_QM, 
                                    start = list(21,4), 
                                    end = list(21,6),
                                    time_dim = 'ftime', ncores = 6)

sprR_fcst$data <- MeanDims(sprR_fcst$data, c('lat', 'lon'))
sprR_hcst_ref$data <- MeanDims(sprR_hcst_ref$data, c('lat', 'lon'))
sprR_obs$data <- MeanDims(sprR_obs$data, c('lat', 'lon'))
sprR_hcst$data <- MeanDims(sprR_hcst$data, c('lat', 'lon'))


metric <- RPSS(sprR_hcst$data, obs = sprR_hcst_ref$data)

terciles <- quantile(sprR_hcst$data, c(0.33, 0.66))
extremes <- quantile(sprR_hcst$data, c(0.10, 0.90))

dim(sprR_fcst$data) <- c(member = 50, sdate = 1)
PlotForecastPDF(sprR_fcst$data,
                tercile.limits = as.vector(terciles),
                extreme.limits = extremes,
                var.name = "sprR (mm)",
                title = "Seasonal forecasts at West Iberian Peninsula",
                fcst.names = paste("Start date: 2022-04-01\n fRPSS:",
                                   round(metric$rpss, 3)),
                obs = as.vector(sprR_obs$data),
                plotfile = "sprR_PlotForecast_csindicators_2022.png")