# Step1: 
library(CSTools)
library(ClimProjDiags)
era5 <- list(name = 'era5',
             path = '/esarchive/scratch/nperez/ERA5/daily_total_prlr_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)
range(era5$data, na.rm = T)
era5$data <- era5$data * 1000 # with or without this line -> same result
era5 <- CST_SplitDim(era5, split_dim = 'sdate', indices = rep(1:12, 26))
slope <- CST_RFSlope(era5, time_dim = c('sdate', 'ftime'), kmin = 5)
print(as.vector(slope))
print(dim(slope))
save(slope,
     file = '/esarchive/scratch/nperez/CSTools_manuscript/v05052021/Slope.RDS',
     version = 2)
#---------------------------------------
# Step2:
library(zeallot)
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')
#obs <- list(name = 'era5', path = '/esarchive/scratch/nperez/ERA5/daily_total_prlr_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 <- CST_SplitDim(exp, split_dim = c('ftime'))
obs <- CST_SplitDim(obs, split_dim = c('ftime'))
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))
print(dim(exp$data))
print(dim(obs$data))
ini <- Sys.time()
exp.qm <- CST_QuantileMapping(exp, obs, method = "QUANT",
                              wet.day = FALSE,
                              sample_dims = c('member', 'sdate', 'ftime'),
                              ncores = 4)
fin <- Sys.time()
print(fin-ini)
save(exp.qm,
     file = '/esarchive/scratch/nperez/CSTools_manuscript/v05052021/ExpQM.RDS',
     version = 2)
# ---------------------------------
# STEP 3: To be run with Internet connexion
#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 = 4)
#save(weight,
#    file = '/esarchive/scratch/nperez/CSTools_manuscript/v05052021/weights.RDS')

# ONLY Plotting on Nord3
load('/esarchive/scratch/nperez/CSTools_manuscript/v05052021/weights.RDS')
library(ragg); library(s2dv)
agg_png("/esarchive/scratch/nperez/CSTools_manuscript/v05052021/WeightsDec.png",
        width = 500, height = 600, units = 'px',res = 144)
PlotEquiMap(weight$data[,,12], lon = weight$lon, lat = weight$lat,
            filled.continents = FALSE,
            intylat = 2, intxlon = 2,
            toptitle = 'December Weights WorldClim2') 
dev.off()
agg_png("/esarchive/scratch/nperez/CSTools_manuscript/v05052021/WeightsNov.png",
        width = 500, height = 600, units = 'px',res = 144)
PlotEquiMap(weight$data[,,11], lon = weight$lon, lat = weight$lat,
            filled.continents = FALSE,
            intylat = 2, intxlon = 2,
            toptitle = 'December Weights WorldClim2')
dev.off()

#----------------------------
# STEP 4:
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')
print("DOWN...")
ini <- Sys.time()
fs <- CST_RainFARM(exp.qm, nf = 4,
                   weights = weights, slope = slope,
                   kmin = 1, nens = 10, verbose = TRUE,
                   time_dim = c("member", "sdate", "ftime"), nprocs = 8,
                   drop_realization = TRUE)
fin <- Sys.time()
print(fin-ini)
newfs <- CST_MergeDims(fs, merge_dims = c("ftime", "monthly"),
                         na.rm = TRUE)

newfs$Dates[[1]] <- exp$Dates[[1]]
CST_SaveExp(newfs, destination = '/esarchive/scratch/nperez/CSTools_manuscript/v05052021/')

print(pryr::mem_used())
# PLOT:
print("PLOTTING")
agg_png("/esarchive/scratch/nperez/CSTools_manuscript/v05052021/EXP_11dec.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/v05052021/EXPQM_11dec.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/v05052021/Down_11dec.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()

