#----------------------------------------------------
# Appendix B – Code used for precipitation downscaling
#-----------------------------------------------------
#!/usr/bin/env Rscript
# R script to downscale precipitation seasonal forecasts with the
#RainFARM method (D'Onofrio et al., 2014; Terzago et al., 2018)
# using functions of the R CSTools package
# 1. Load seasonal forecasts & ERA5 daily precip data, both at 1°
#spatial resolution
# 2. Perform quantile mapping to correct ECDF on a monthly basis.
#Reference ERA5 1° res
# 3. Get Monthly Slopes from ERA5 at 25 km
# 4. Estimate orographic weights (monthly basis) from ERA5 at 25
km
# 5. Downscale month by month from 1° to 0.25°
# Use: Rscript
#/work/users/silvia/medscope/downscaling_med/rainfarm-downscale.R
#model start_date_MMDD
# Example: Rscript
#/work/users/silvia/medscope/downscaling_med/rainfarm-downscale.R
#ecmwfS5 0501
#
# Pre-requisite:
# Get ERA5 data:
#./work/users/silvia/medscope/downscaling_med/era5_data/get_era5.sh
# Get model data:
#./work/users/silvia/medscope/downscaling_med/seasfc_data/download_
#med_mmday.shi
#-----------------------------------------------------
library(CSTools)
library(zeallot)
library(fields)
library(raster)
model <- 'ecmwfS5'
mmdd <- '0501'
#wdir <- "/work/users/silvia/medscope/downscaling_med/"

#outpath <- paste0(wdir, "output_qm/", model, "/")
#outdir <- paste0(outpath, "experiment/prlr/")
#downdir <- "/work/users/silvia/medscope/downscaling_med/downscaled/"
#system(paste("mkdir", outpath))
era5clim <- "/work/users/silvia/medscope/CSTools/era5_tp_mon_latlon.nc"
grid <- "/work/users/silvia/medscope/CSTools/clim.grd"
infile <- paste0("/work/users/silvia/medscope/downscaling_med/seasfc_data/", model, "/$VAR_NAME$_$START_DATE$.nc")
reffile <- '/work/users/silvia/medscope/downscaling_med/era5_data/era5_$VAR_NAME$_$START_DATE$.nc'
# Spectral slopes (monthly scale) derived from ERA5 0.25°
# resolution over the domain 2.5/18.25/37.5/53.25
# Slopes calculated from wavenumber >= 7
obs <- list(name = 'era5',
            path = '/esarchive/recon/ecmwf/era5/$STORE_FREQ$_mean/$VAR_NAME$_*/$VAR_NAME$_$YEAR$$MONTH$.nc')
years <- paste0(1993:2018, '0101')
obs <- CST_Load(var = 'prlr',
                exp = NULL,
                obs = list(obs),
                sdates = years, nmember = 25, leadtimemax = 12,
                storefreq = "daily", sampleperiod = 1,
                latmin = 0, latmax = 79.5, lonmin = -19, lonmax = 60.5,
                output = 'lonlat')
slope <- CST_RFSlope(obs, kmin = 7, time_dim = 'sdate')
slope.mon <- c(2.45, 2.43, 2.43, 2.36, 2.22, 2.05, 1.98, 2.06, 2.37, 2.47, 2.52, 2.46)
mon <- c(5:11)
ltn <- c(1, 32, 62, 93, 124, 154, 185)
ltx <- c(31, 61, 92, 123, 153, 184, 214)
down_fact = 4
nrealiz = 5
#starting dates
StartDates <- paste0(1993:2018, '1101')
# Load them separately to have original resolution
print("Load Seasonal Forecasts for the Bias Adjustment by Quantile Mapping")
#exp <- list(name = 'ecmwfS5',
#            path = "/work/users/silvia/medscope/downscaling_med/seasfc_data/ecmwfS5/$VAR_NAME$_$START_DATE$.nc")
exp <- list(name = 'ECMWFs5',
            path = '/esarchive/exp/ecmwf/system5c3s/$STORE_FREQ$_mean/$VAR_NAME$_s0-24h/$VAR_NAME$_$START_DATE$.nc')
exp <- CST_Load(var = 'prlr',
                exp = list(exp),
                obs = NULL,
                sdates = StartDates, nmember = 25, leadtimemax = 214,
                storefreq = "daily", sampleperiod = 1,
                latmin = 0, latmax = 79.5, lonmin = -19, lonmax = 60.5,
                output = 'lonlat')
#obs <- list(name = 'era5',
#            path = '/work/users/silvia/medscope/downscaling_med/era5_data/era5_$VAR_NAME$_$START_DATE$.nc')
obs <- list(name = 'era5',
            path = '/esarchive/recon/ecmwf/era5/$STORE_FREQ$_mean/$VAR_NAME$_*/$VAR_NAME$_$YEAR$$MONTH$.nc')
obs <- CST_Load(var = 'prlr',
                exp = NULL,
                obs = list(obs),
                sdates = StartDates, nmember = 1, leadtimemax = 214,    
                storefreq = "daily", sampleperiod = 1,
                latmin = 0, latmax = 79.5, lonmin = -19, lonmax = 60.5,
                output = 'lonlat')
# SplitDim
print("Splitting ftime by month ...")
expm <- CST_SplitDim(exp, split_dim = c('ftime'))
obsm <- CST_SplitDim(obs, split_dim = c('ftime'))
obsm$data <- obsm$data[,,,,,,1:7]
# CST_QuantileMapping by month:
# sample_dims: a character vector indicating the dimensions that
# can be used as sample for the same distribution
# Please note that wet.day=T would be optimal but not enough nonzero
# data in a month to work
exp.qm <- CST_QuantileMapping(expm,obsm, method = "QUANT", wet.day = FALSE,
                              sample_dims = c('member', 'sdate', 'ftime'))
# MergeDims
print("Restoring original dimensions ...")
exp.qmm <- CST_MergeDims(exp.qm, merge_dims = c("ftime", "monthly"), na.rm = TRUE)
# Adjust longitude to be monotonic [-180;180]
lon <- as.numeric(exp.qmm$lon)
exp.qmm$lon <- lon * (lon <= 180) + (lon - 360) * (lon > 180)
lat = exp.qmm$lat
lon = exp.qmm$lon
l <- length(lon)
lf <- l * down_fact
#print (paste("Save Bias-Adjusted data to", outdir)) 
#CST_SaveExp(exp.qmm, outpath)
#rm(exp, expm, exp.qm, exp.qmm)
##########################################
# Downscaling by month (slope by month)
print("Subsetting data by month ... ")
infile <- paste0(outdir, "$VAR_NAME$_$START_DATE$.nc")
exp <- list(name = 'experiment', path = infile)
for (i in 1:length(mon)){
  # Check month
  mm <- sprintf("%02d", mon[i])
  print(mm)
  system(paste0("mkdir ", downdir, "mon_", mm))
  # Use slopes from ERA5
  slo <- slope.mon[mon[i]]
  # Get weights from ERA5
  print("Get weights from ERA5")
  climmon <- paste0("clim_" ,mm,".nc")
  system(paste("rm", climmon))
  cmd0 <- paste0("cdo remapnn,", grid, " -selmon,", mm, " ", era5clim, " ", climmon)
  system(cmd0)
  w_clim_tmp <- CST_RFWeights(climmon, lat = lat, lon = lon, nf = down_fact)
  # Get rid of NaN in the weights
  r <- raster(nrow = lf, ncol = lf)
  r[] <- w_clim_tmp
  base <- raster(nrow = lf, ncol = lf)
  if (length(table(is.na(as.matrix(r)))) > 1) {
    while (table(is.na(as.matrix(r)))[2] > 0) {
      r <- resample(r, base, method = 'bilinear')
      if (length(table(is.na(as.matrix(r)))) == 1) break
    }
  }
  w_clim <- as.matrix(r)
  # Get seasonal forecast for month mm, loop on start-date
  print("Get seasonal forecast for month mm, loop on startdate")
  for (j in 1:length(startDates)){
    if ((mmdd == "1101") && (startDates == 1995 | startDates == 1999 | startDates == 2003 |
                             startDates == 2007 |startDates == 2011 | startDates == 2015)) {
      print ("Anno Bisestoxxx")
      ltn <- c(1, 31, 62, 93, 122, 153, 183)
      ltx <- c(30, 61, 92, 121, 152, 182, 213)
    }
    print(startDates[j])
    ls %<-% CST_Load(var = 'prlr',
                     exp = list(exp),
                     obs = NULL,
                     sdates = startDates[j], nmember = 25,
                     leadtimemin = ltn[i], leadtimemax = ltx[i],
                     storefreq = "daily", sampleperiod = 1,
                     latmin = 0, latmax = 79.5, lonmin = -19, lonmax = 60.5,
                     output = 'lonlat',nprocs=6)
    lon <- ls$lon
    ls$lon <- lon * (lon <= 180) + (lon - 360) * (lon > 180)
    # Orographic downscaling
    print("Orographic downscaling")
    for (k in 1:nrealiz) {
      k2 <- sprintf("%03d", k)
      fs <- CST_RainFARM(ls, nf = down_fact,
                         weights = w_clim, slope = slo, kmin = 1, nens = 1, verbose = TRUE,
                         time_dim = c("member", "sdate", "ftime"), nprocs = 6,
                         drop_realization = TRUE)
      # Save output
      print(paste0("Saving output to: ",downdir, "mon_", mm, "experiment/prlr/"))
      CST_SaveExp(fs, paste0(downdir, "mon_", mm))
      filetmp=paste0(downdir, "mon_", mm, "/experiment/prlr/prlr_", startDates[j], ".nc")
      fileout=paste0(downdir, "mon_", mm,"/experiment/prlr/prlr_", startDates[j], "_rainfarm_", k2 ,".nc")
      system(paste("mv", filetmp, fileout))
      rm(fs)
    }
    rm (ls)
  }
}
# Mergetime
print (paste0("Merge monthly downscaled data by starting date, then write output to: ", downdir, model, "/"))
system(paste0("mkdir ", downdir, model, "/"))
for (j in 1:length(startDates)){
  for (k in 1:nrealiz){
    k2 <- sprintf("%03d", k)
    nfiles <- paste0(downdir, "mon_*/experiment/prlr/prlr_", startDates[j], "_rainfarm_", k2
,".nc")
    downfile <- paste0(downdir, model, "/prlr_",
    startDates[j], "_rainfarm_", k2 , ".nc")
    system(paste("cdo mergetime", nfiles, downfile))
  }
}
# End downscaling

