rm(list=ls())

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

source("https://earth.bsc.es/gitlab/nperez/Flor/-/raw/master/CST_Load_devel_from_s2dv.R")

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_s2dv(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_s2dv(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'


source("https://earth.bsc.es/gitlab/external/cstools/-/raw/develop-QuantMap/R/CST_QuantileMapping.R")

library(qmap)
library(multiApply)
library(abind)

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)

PlotEquiMap(fcst$data[1,1,1,2,,], lon = fcst$lon, lat = fcst$lat, filled.c = F)
dev.new()
PlotEquiMap(fcst_QM$data[1,1,2,,], lon = fcst$lon, lat = fcst$lat, filled.c = F)

fcst_QM$data <- InsertDim(fcst_QM$data, posdim = 1, lendim = 1, name = 'sdate')

source("https://earth.bsc.es/gitlab/es/csindicators/-/raw/master/R/SelectPeriodOnData.R")


sprR_fcst <- CST_PeriodAccumulation(fcst_QM,
                                    start = list(21,4), 
                                    end = list(21,6), na.rm = FALSE,
                                    time_dim = 'ftime', ncores = 6)
sprR_hcst <- CST_PeriodAccumulation(hcst, 
                                    start = list(21,4), 
                                    end = list(21,6),
                                    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_QM <- 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$data <- MeanDims(sprR_hcst$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_QM$data <- MeanDims(sprR_hcst_QM$data, c('lat', 'lon'))



source("https://earth.bsc.es/gitlab/es/s2dv/-/raw/master/R/RPSS.R")
source("https://earth.bsc.es/gitlab/es/s2dv/-/raw/master/R/RPS.R")
source("https://earth.bsc.es/gitlab/es/s2dv/-/raw/master/R/RandomWalkTest.R")

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

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

dim(sprR_fcst$data) <- c(members = 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"))

data <- plyr::adply(sprR_hcst$data, c(1,2))
 
boxplot(data[,-c(1,2)], xlab = "Time (year)", ylab = "sprR (mm)", names = 1993:2016, 
        col = 'lightgreen')
points(1:24, sprR_hcst_ref$data, col = 'red', pch = 16)


plot(1:24, rep(NA, 24), xlim = c(0, 25), ylim = c(0, 150), 
     ylab = 'sprR (mm)', xlab = "Start Dates (years)", bty = 'n')

polygon(x = c(0.75, 1.25, 0.75, 1.25), 
        y = c(rep(min(sprR_hcst$data[1,,1]), 2), 
              rep(quantile(sprR_hcst$data[1,,1],0.33), 2)))