#!/usr/bin/env Rscript
rm(list=ls()); gc();
library(CSTools)
source("/esarchive/scratch/nperez/git/cstools/R/CST_RFSlope.R")
source("/esarchive/scratch/nperez/git/cstools/R/CST_RainFARM.R")
source("/esarchive/scratch/nperez/git/cstools/R/CST_SplitDim.R")
source("/esarchive/scratch/nperez/git/cstools/R/CST_RFWeights.R")
source("/esarchive/scratch/nperez/git/cstools/R/CST_QuantileMapping.R")
source("/esarchive/scratch/nperez/git/cstools/R/zzz.R")
source("/esarchive/scratch/nperez/git/cstools/R/CST_SaveExp.R")
library(abind)
library(s2dverification)
library(qmap)
library(multiApply)
library(ncdf4); library(s2dv)
library(rainfarmr)
library(ClimProjDiags)
library(zeallot)
#----------
# STEP 1:
print("STEP 1")
era5 <- list(name = 'era5', path = '/esarchive/recon/ecmwf/era5/$STORE_FREQ$_mean/$VAR_NAME$_f1h-r1440x721cds/$VAR_NAME$_$YEAR$$MONTH$.nc')
years <- unlist(lapply(1993:2018, function(x){paste0(x, sprintf("%02d",1:12), '01')}))
era5 <- CST_Load(var = 'prlr', exp = list(era5), sdates = years, nmember = 1,
                               storefreq = "daily", sampleperiod = 1, 
                               latmin = 37.5, latmax = 53.25, lonmin = 2.5, lonmax = 18.25, 
                               output = 'lonlat', nprocs = 1)
era5$data <- era5$data * 24 * 3600 * 1000
era5 <- CST_SplitDim(era5, split_dim = 'sdate', indices = rep(1:12, 26))
slope <- CST_RFSlope(era5, time_dim = c('sdate', 'ftime'), kmin = 5)

print(slope)
print(dim(slope))
#----------
# STEP 2:
print("STEP 2")
StartDates <- paste0(1993:2018, '1101')
exp <- list(name = 'ecmwfS5', path = "/esarchive/exp/ecmwf/system5c3s/$STORE_FREQ$_mean/$VAR_NAME$_s0-24h/$VAR_NAME$_$START_DATE$.nc")
obs <- list(name = 'era5', path = '/esarchive/recon/ecmwf/era5/$STORE_FREQ$_mean/$VAR_NAME$_f1h-r1440x721cds/$VAR_NAME$_$YEAR$$MONTH$.nc')
c(exp, obs) %<-% CST_Load(var = 'prlr', exp = list(exp), obs = list(obs),
                          sdates = StartDates, nmember = 25,
                          storefreq = "daily", sampleperiod = 1,
                          latmin = 42, latmax = 49, lonmin = 4, lonmax = 11,
                          output = 'lonlat', nprocs = 1)

exp$data <- exp$data * 24 * 3600 * 1000
obs$data <- obs$data * 24 * 3600 * 1000
exp$data[which(exp$data < 0)] <- 0

print(range(obs$data, na.rm = TRUE))
print(range(exp$data, na.rm = TRUE))

exp <- CST_SplitDim(exp, split_dim = c('ftime'), insert_ftime = 1)
obs <- CST_SplitDim(obs, split_dim = c('ftime'), insert_ftime = 1)

exp.qm <- CST_QuantileMapping(exp, obs, method = "QUANT",
                                     wet.day = FALSE,
                                     sample_dims = c('member', 'sdate', 'ftime'), ncores = 14)
save(obs, file = '/esarchive/scratch/nperez/CSTools_manuscript/Dataset_Obs.RDS', version = 2)
save(exp, file = '/esarchive/scratch/nperez/CSTools_manuscript/Dataset_Exp.RDS', version = 2)
save(exp.qm, file = '/esarchive/scratch/nperez/CSTools_manuscript/Dataset_ExpQM.RDS', version = 2)

#----------------------------------
# STEP 3: to run with connection to internet
# WorldClim data to s2dv_cube
#library(CSTools); library(raster); library(s2dv)
#worldclim <- getData("worldclim", var = "prec", res = 0.5, lon = 5, lat = 45)
#wc_month <- lapply(1:12, FUN = function(x) {
#                   res <- crop(worldclim[[x]],
#                               extent(3.5, 11.5, 41.5, 49.5))
#                   res <- as.array(res)
#                   names(dim(res)) <- c('lat', 'lon', 'month')
#                   return(res)
#                   })
#xy <- xyFromCell(crop(worldclim[[1]], extent(3.5, 11.5, 41.5, 49.5)),
#                 1:length(crop(worldclim[[1]], extent(3.5, 11.5, 41.5, 49.5))))
#lons <- unique(xy[,1])
#lats <- unique(xy[,2])
#wc_month <- unlist(wc_month)
#dim(wc_month) <- c(lat = length(lats), lon = length(lons), monthly = 12)
#wc_month <- Reorder(wc_month, c('lon', 'lat', 'monthly'))
#wc_month <- s2dv_cube(data = wc_month, lon = lons, lat = lats,
#                       Datasets = 'WorldClim')
#weight <- CST_RFWeights(wc_month, lon = 4:11, lat = 49:42, nf = 100) 
#weight <- CST_RFWeights(wc_month, lon = exp$lon, lat = exp$lat, nf = 100) # if all step 2 is run
#save(weight, file = '/esarchive/scratch/nperez/CSTools_manuscript/Dataset_weight.RDS', version = 2)

#----------
# STEP 4:
print("STEP 4")
load('/esarchive/scratch/nperez/CSTools_manuscript/Dataset_weight.RDS')
weights <- Subset(weight$data, along = 'monthly', indices = c(11,12,1:1:6))
slope <- Subset(slope, along = 'monthly', indices = c(11,12,1:1:6), drop = 'non-selected')
exp.qm2 <- exp.qm
fs <- NULL
for (i in 1:2) {
print(paste("LOOP", i))
exp.qm2$data <- Subset(exp.qm$data, along = 'sdate', indices = i)
  
fs2 <- CST_RainFARM(exp.qm2, nf = 100,
                   weights = weights, slope = slope,
                   kmin = 1, nens = 10, verbose = TRUE,
                   time_dim = c("member", "ftime"), nprocs = 16,
                   drop_realization = TRUE)
pos <- which(names(dim(fs2$data)) == 'sdate')
fs <- abind(fs, fs2$data, along = pos) 
}
save(fs, file = '/esarchive/scratch/nperez/CSTools_manuscript/Dataset_FS.RDS', version = 2)
fs_merged <- CST_MergeDims(fs, merge_dims = c("ftime", "monthly"), na.rm = TRUE)
 
fs_merged$Dates[[1]] <- exp$Dates[[1]]
CST_SaveExp(fs, destination = '/esarchive/scratch/nperez/CSTools_manuscript/v20201201/')

#----------
# STEP 5: Plotting
print("STEP 5")
library(ragg)
library(ragg);library(s2dv)
agg_png("/esarchive/scratch/nperez/CSTools_manuscript/v20201201/Single1.png",
        width = 500, height = 600, units = 'px',res = 144)
PlotEquiMap(exp$data[1,1,1,11,,,2],lon = exp$lon, lat = exp$lat,
            filled.continents = FALSE, bar_limits = c(0,40),
           intylat = 2, intxlon = 2, title_scale = 0.7,
           toptitle = 'ECMWF-S5C3S', units = 'precipitation (mm)')
dev.off()
agg_png("/esarchive/scratch/nperez/CSTools_manuscript/v20201201/Single2.png",
        width = 500, height = 600, units = 'px',res = 144)
PlotEquiMap(exp.qm$data[1,1,1,11,,,2],lon = exp$lon, lat = exp$lat,
            filled.continents = FALSE, bar_limits = c(0,40),
           intylat = 2, intxlon = 2, title_scale = 0.7,
           toptitle = 'Bias Corrected', units = 'precipitation (mm)')
dev.off()
agg_png("/esarchive/scratch/nperez/CSTools_manuscript/v20201201/Single3.png",
        width = 500, height = 600, units = 'px',res = 144)
PlotEquiMap(fs$data[1,1,1,11,,,2],lon = fs$lon, lat = fs$lat,
            filled.continents = FALSE, bar_limits = c(0,40),
           intylat = 2, intxlon = 2, title_scale = 0.7,
           toptitle = 'Downsacaled', units = 'precipitation (mm)')
dev.off()
agg_png("/esarchive/scratch/nperez/CSTools_manuscript/v20201201/Single4.png",
        width = 500, height = 600, units = 'px',res = 144)
PlotEquiMap(weight$data[,,2],lon = weight$lon, lat = weight$lat,
            filled.continents = FALSE,
           intylat = 2, intxlon = 2, title_scale = 0.7,
           toptitle = 'Weights', units = 'weights(adim)')
dev.off()
agg_png("/esarchive/scratch/nperez/CSTools_manuscript/v20201201/Single5.png",
        width = 500, height = 600, units = 'px',res = 144)
plot(1:12, slope[3:12,1:2] type = 'b', main = 'slope')
dev.off()

# Plot ForecastPDF
obsts <- MeanDims(obs$data, c('lat', 'lon'), na.rm = T)
print(quantile(obsts, c(0.1,0.3,0.6,0.9), na.rm = T))
expts <- MeanDims(exp$data, c('lat', 'lon'), na.rm = T)
exp.qmts <- MeanDims(exp.qm$data,  c('lat', 'lon'), na.rm = T)
empty <- array(NA, c(dataset = 2, member = 225, sdate = 26, ftime = 31, monthly = 8))
data <- abind(expts, exp.qmts, along = 1)
names(dim(data)) <- names(dim(expts))
data <- abind(data, empty, along = 2)
names(dim(data)) <- names(dim(expts))
fsts <- MeanDims(fs, c('lat', 'lon'), na.rm = T)
data <- abind(data, fsts, along = 1)
names(dim(data)) <- c('dataset', 'members', 'sdate', 'ftime', 'monthly')
agg_png("/esarchive/scratch/nperez/CSTools_manuscript/v20201201/FiguresPDF_11DecemberAll2.png",
        width = 750, height = 650, units = 'px', res = 144)
PlotForecastPDF(data[,,1,11,2], tercile.limits = c(0.58, 2.4), obs = obsts[,,1,11,2], 
    extreme.limits = c(0.06, 7.26), color.set = 'hydro', add.ensmemb = 'no',
    var.name = "Precipitation (mm)", title = "Forecasts issued on Nov 1993 for 11th December 1993",
    fcst.names = c("ECMWFS5C3S", "Bias Corrected", "Downscaled"))
dev.off()


