#!/usr/bin/env Rscript
rm(list=ls()); gc();
library(CSTools)
library(zeallot)
library(ragg)
library(s2dv)
# Already computed weight and slopes:
load('/esarchive/scratch/nperez/CSTools_manuscript/Dataset_weight.RDS')
weights <- ClimProjDiags::Subset(weight$data, along = 'monthly',
                                 indices = c(11,12,1:1:6))
slope <- c(2.45, 2.43, 2.44, 2.36, 2.20, 2.05, 1.98, 2.06, 2.37, 2.47, 2.52, 2.46)
dim(slope) <- c(dataset = 1, monthly = 12)
slope <- ClimProjDiags::Subset(slope, along = 'monthly',
                               indices = c(11,12,1:1:6), drop = 'non-selected')
agg_png("/esarchive/scratch/nperez/CSTools_manuscript/v30032021/weights.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()

# Load version that doesn't check the first time step in files:
#source("/esarchive/scratch/nperez/git/Flor/CSTools_Manuscript/Load.R")
#source("/esarchive/scratch/nperez/git/Flor/CSTools_Manuscript/Utils.R")

StartDates <- paste0(1993:2018, '1101')

exp_path <- list(name = 'ecmwfS5', path = "/esarchive/exp/ecmwf/system5c3s/$STORE_FREQ$/$VAR_NAME$_s0-24h/$VAR_NAME$_$START_DATE$.nc")
        
obs_path <- 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_path), obs = list(obs_path),
                          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[exp$data < 0] <- 0

agg_png("/esarchive/scratch/nperez/CSTools_manuscript/v30032021/exp_raw_11Dec1993.png",
        width = 500, height = 600, units = 'px',res = 144)
PlotEquiMap(exp$data[1, 1, 1, 41,,], lon = exp$lon, lat = exp$lat,
            filled.continents = F, 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/v30032021/obs_raw_11Dec1993.png",
        width = 500, height = 600, units = 'px',res = 144)
PlotEquiMap(obs$data[1, 1, 1, 41,,], lon = obs$lon, lat = obs$lat,
            filled.continents = F,  bar_limits = c(0,40),
            intylat = 2, intxlon = 2, title_scale = 0.7,
           toptitle = 'ERA5', units = 'precipitation (mm)')
dev.off()

exp <- CST_SplitDim(exp, split_dim = c('ftime'))
obs <- CST_SplitDim(obs, split_dim = c('ftime'))
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/v30032021/Dataset_Obs.RDS',
     version = 2)
save(exp,
     file = '/esarchive/scratch/nperez/CSTools_manuscript/v30032021/Dataset_Exp.RDS',
     version = 2)
save(exp.qm,
     file = '/esarchive/scratch/nperez/CSTools_manuscript/v30032021/Dataset_ExpQM.RDS',
     version = 2)
agg_png("/esarchive/scratch/nperez/CSTools_manuscript/v30032021/BC_QM_11Dec1993.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()

QM <- CST_MergeDims(exp.qm, na.rm = TRUE)
CST_SaveExp(QM, 
destination = "/esarchive/scratch/nperez/CSTools_manuscript/v30032021/Corrected/")

fs <- NULL
exp.qm2 <- exp.qm
#for (i in 1:2) {
i = 1
print(paste("LOOP", i))
exp.qm2$data <- ClimProjDiags::Subset(exp.qm$data, along = 'sdate', indices = i)
exp.qm3 <- exp.qm2
 for (j in 1:8) {
   exp.qm3$data <- ClimProjDiags::Subset(exp.qm$data, along = 'monthly', indices = j)  
 fs2 <- CST_RainFARM(exp.qm2, nf = 100,
                   weights = weights[,,j], slope = slope,
                   kmin = 1, nens = 10, verbose = TRUE,
                   time_dim = c("member", "ftime"), nprocs = 16,
                   drop_realization = TRUE)

CST_SaveExp(fs2,
            destination = "/esarchive/scratch/nperez/CSTools_manuscript/v30032021/")

#pos <- which(names(dim(fs2$data)) == 'sdate')
#fs <- abind(fs, fs2$data, along = pos) 
 if (j== 2) { 
  print(dim(fs2$data))
  agg_png("/esarchive/scratch/nperez/CSTools_manuscript/v30032021/Down_11Dec1993.png",
        width = 500, height = 600, units = 'px',res = 144)
  PlotEquiMap(fs2$data[1,1,1,11,,,1],lon = fs2$lon, lat = fs2$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()
 }
gc()
}

stop("End of downscaling")

