#----------------------------------------------------
# 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(devtools)
library(zeallot)
library(fields)
library(raster)
load_all()
args = commandArgs(trailingOnly=TRUE)
# test if there is at least one argument: if not, return an error
if (length(args) < 2) {
  stop("Two arguments must be supplied: model , startDate MMDD", call.=FALSE)
} else {
  print(paste("Model: ", args[1]))
  print(paste("Starting date MMDD: ", args[2]))
}
# end test
######################
# get or set variables
model <- args[1]
mmdd <- args[2]
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
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)
if (mmdd == "0501") {
  mon <- c(5:11)
  ltn <- c(1, 32, 62, 93, 124, 154, 185)
  ltx <- c(31, 61, 92, 123, 153, 184, 214)
} else if (mmdd == "1101") {
  mon <- c(11, 12, 1:5)
  ltn <- c(1, 31, 62, 93, 121, 152, 182)
  ltx <- c(30, 61, 92, 120, 151, 181, 212)
}
var = 'prlr'
nens = 25
lonn = -19
lonx = 60.5
latn = 0
latx = 79.5
down_fact = 4
nrealiz = 5
y1 = 1993
y2 = 2018
yrs <- y1:y2
nsd <- length(yrs)
sd <- vector()
for (i in 1:nsd) {
  sd[i] <- paste0(yrs[i], mmdd)}
  startDates <- sd
  # end get or set variables
  ##################################################################
  # Load Seasonal Forecasts for the Bias Adjustment by Quantile Mapping
  print("Load Seasonal Forecasts for the Bias Adjustment by Quantile Mapping")
  exp <- list(name = 'experiment', path = infile)
  obs <- list(name = 'observation',path = reffile)
  exp %<-% CST_Load(var = 'prlr',
                    exp = list(exp)
                    obs = NULL,
                    sdates = startDates, nmember = nens, 
                    storefreq = "daily", sampleperiod = 1,
                    latmin = latn, latmax = latx, lonmin = lonn, lonmax = lonx,
                    output = 'lonlat', nprocs = 6)
  obs %<-% CST_Load(var = 'prlr',
                    exp = list(obs),
                    obs = NULL, 
                    sdates = startDates, nmember = 1,
                    storefreq = "daily", sampleperiod = 1,
                    latmin = latn, latmax = latx, lonmin = lonn, lonmax = lonx,
                    output = 'lonlat', nprocs = 6)
  dim(exp$data)
  dim(obs$data)
  # 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]
  dim(expm$data)
  dim(obsm$data)
  # 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 = F,
                                sample_dims = c('member', 'sdate', 'ftime'))
  dim(exp.qm$data)
  # 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('prlr', exp=list(exp), obs=NULL,
                       sdates =startDates[j], nmember = nens, leadtimemin = ltn[i],
                       leadtimemax=ltx[i],
                       storefreq = "daily", sampleperiod = 1,
                       latmin =latn, latmax = latx, lonmin = lonn, lonmax = lonx,
                       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=T)
        # 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

