#!/usr/bin/env Rscript
rm(list=ls()); gc();

### Creation date: 3rd June 2021
# Author: N. Pérez-Zanón
# Refer CSTools package and manuscript when using it.
# ----------------------------------------
# Wind speed bias adjustment for assessment of an extreme event
# The System5-ECMWF downloaded from C3S seasonal forecasts initialized 
# in December 2017, January 2018 and February 2018 
# This code includes the bias adjustent and the results visualization 
# ----------------------------------------

library(CSTools)
library(s2dv)
library(ragg)
library(multiApply)
output_dir <- "/esarchive/scratch/nperez/ArticXChange/"


exp_path <- list(name = "ECMWFS5", 
                 path = "/esarchive/exp/ecmwf/system5c3s/$STORE_FREQ$_mean/$VAR_NAME$_f6h/$VAR_NAME$_$START_DATE$.nc")
obs_path <- list(name = "ERA5",
                 path = "/esarchive/recon/ecmwf/era5/$STORE_FREQ$_mean/$VAR_NAME$_f1h/$VAR_NAME$_$YEAR$$MONTH$.nc")
#source("/esarchive/scratch/nperez/git/cstools/R/CST_BiasCorrection.R")
#library(multiApply)
  # Target months March (3)
  # Assess forecast from 1 to 3 months in advance
months_in_advance <- c('05', '04', '03')
year <- 2020
target_m <- '06'
areave <- 'areave' # 'areave' or 'lonlat' #best option areave
if (areave == 'lonlat') {
  grid <- 'r360x181'
} else {
  grid <- NULL
}
tas_fsct_BC <- list()
tas_ref_terciles <- NULL
tas_ref_extremes <- NULL
tas_hcst_terciles <- NULL
tas_hcst_extremes <- NULL

# For each month in advance:
for (mm in 1:3) {
  print("Three initializations:")
  print(mm)
  print(paste('Initialization', months_in_advance[mm]))
  # Generate the start dates of hindcast period
  print(paste('Hindcast period until', year - 1))
  hcst_sdates <- paste0(1993:(year - 1), months_in_advance[mm], '01') 
  tas_hcst <- CSTools::CST_Load(var = 'tas', exp = list(exp_path),
                        sdates = hcst_sdates, nmember = 25, 
                        leadtimemin = mm + 1, leadtimemax = mm + 1,
                        storefreq = "monthly", sampleperiod = 1,
                        latmin = 65, latmax = 71, lonmin = 20, lonmax = 31,
                        output = areave, nprocs = 1)
  str(tas_hcst$Dates)
  dim(tas_hcst$data)
  fcst_sdates <- paste0(year, months_in_advance[mm], '01')
  tas_fcst <- CSTools::CST_Load(var = 'tas', exp = list(exp_path),
                        sdates = fcst_sdates, nmember = 25,
                        leadtimemin = mm + 1, leadtimemax = mm + 1,
                        storefreq = "monthly", sampleperiod = 1,
                        latmin = 65, latmax = 71, lonmin = 20, lonmax = 31,
                        output = areave, nprocs = 1)
  str(tas_fcst$Dates)
  dim(tas_fcst$data)
  
  tas_ref <- CSTools::CST_Load(var = 'tas', obs = list(obs_path),
                       sdates = hcst_sdates, nmember = 1,
                       leadtimemin = mm + 1, leadtimemax = mm + 1,
                       storefreq = "monthly", sampleperiod = 1,
                       latmin = 65, latmax = 71, lonmin = 20, lonmax = 31,
                       output = areave, nprocs = 1,
                       grid = grid)
  str(tas_ref$Dates)
  dim(tas_ref$data)
  print(tas_ref$Dates$start)
  if (areave == 'lonlat') {                                  
    tas_ref_terciles <- rbind(tas_ref_terciles,
                            quantile(MeanDims(tas_ref$data, c('lat', 'lon')),
                                     c(0.3, 0.6)))
    tas_ref_extremes <- rbind(tas_ref_extremes,
                            quantile(MeanDims(tas_ref$data, c('lat', 'lon')),
                            c(0.1, 0.9))) 
    tas_hcst_terciles <- rbind(tas_hcst_terciles,
                             quantile(MeanDims(tas_hcst$data, c('lat', 'lon')),
                                      c(0.3, 0.6)))
    tas_hcst_extremes <- rbind(tas_hcst_extremes,
                            quantile(MeanDims(tas_hcst$data, c('lat', 'lon')),
                            c(0.1, 0.9)))
  } else {
    tas_ref_terciles <- rbind(tas_ref_terciles,
                            quantile(tas_ref$data, c(0.3, 0.6)))
    tas_ref_extremes <- rbind(tas_ref_extremes,
                            quantile(tas_ref$data, c(0.1, 0.9))) 
    tas_hcst_terciles <- rbind(tas_hcst_terciles,
                             quantile(tas_hcst$data, c(0.3, 0.6)))
    tas_hcst_extremes <- rbind(tas_hcst_extremes,
                            quantile(tas_hcst$data, c(0.1, 0.9)))
  } 

 
  tas_fsct <- CST_BiasCorrection(exp = tas_hcst,
                                 obs = tas_ref, 
                                 exp_cor = tas_fcst)
  tas_fsct_BC[[mm]] <- tas_fsct
}
  tas_obs_areave <- CSTools::CST_Load(var = 'tas', obs = list(obs_path),
                        sdates = paste0(year, target_m, '01'), nmember = 1,
                        leadtimemin = 1, leadtimemax = 1,
                        storefreq = "monthly", sampleperiod = 1,
                        latmin = 65, latmax = 71, lonmin = 20, lonmax = 31,
                        output = 'areave', nprocs = 1)

if (areave == 'lonlat') {
  print('lonlat')
  visual <- data.frame(dec = as.vector(MeanDims(tas_fsct_BC[[3]]$data, 
                                     c('lat', 'lon'))) - 273.15,
                       jan = as.vector(MeanDims(tas_fsct_BC[[2]]$data,
                                     c('lat', 'lon'))) - 273.15,
                       feb = as.vector(MeanDims(tas_fsct_BC[[1]]$data,
                                     c('lat', 'lon'))) - 273.15)

  print("IS DATA LOADED")
# Case 1: terciles and extremes are computed with the obs
#         bias correction point by point and later aggregation
  agg_png(paste0(output_dir, "Tas_PlotForecast_1.png"),
    width = 1000, height = 1000, units = 'px',res = 150)
  CSTools::PlotForecastPDF(visual, 
                tercile.limits = tas_ref_terciles - 273.15,
                extreme.limits = tas_ref_extremes - 273.15,
                var.name = "Air temperature at surface",
                title = "Bias Corrected forecasts at Lappland (1)", 
                fcst.names = month.name[rev(as.numeric(months_in_advance))], 
                obs = as.vector(tas_obs_areave$data) - 273.15)
  dev.off()

# Case 2: terciles and extremes are computed with the hcst
#         bias correction point by point and later aggregation
  agg_png(paste0(output_dir, "Tas_PlotForecast_2.png"),
    width = 1000, height = 1000, units = 'px',res = 150)
  CSTools::PlotForecastPDF(visual,
                tercile.limits = tas_hcst_terciles[3:1,] - 273.15,
                extreme.limits = tas_hcst_extremes[3:1,] - 273.15,
                var.name = "Air temperature at surface",
                title = "Bias Corrected forecasts at Lappland (2)",
                fcst.names = month.name[rev(as.numeric(months_in_advance))],
                obs = as.vector(tas_obs_areave$data) - 273.15)
  dev.off()
} else {
print("areave")
 visual <- data.frame(dec = as.vector(tas_fsct_BC[[3]]$data) - 273.15,
                       jan = as.vector(tas_fsct_BC[[2]]$data) - 273.15,
                       feb = as.vector(tas_fsct_BC[[1]]$data) - 273.15)
# Case 3: terciles and extremes are computed with the obs
#         bias correction is computed on the areave data
  agg_png(paste0(output_dir, "Tas_PlotForecast_3.png"),
    width = 1000, height = 1000, units = 'px',res = 150)
  CSTools::PlotForecastPDF(visual,
                tercile.limits = tas_ref_terciles[3:1,] - 273.15,
                extreme.limits = tas_ref_extremes[3:1,] - 273.15,
                var.name = "Air temperature at surface",
                title = "Bias Corrected forecasts at Lappland (3)",
                fcst.names = month.name[rev(as.numeric(months_in_advance))],
                obs = as.vector(tas_obs_areave$data) - 273.15)
  dev.off()
# Case 4: terciles and extremes are computed with the hcst
#         bias correction is computed on the areave data
  agg_png(paste0(output_dir, "ERA5_Tas_PlotForecast_4_",year, target_m, ".png"),
    width = 1000, height = 1000, units = 'px',res = 150)
  CSTools::PlotForecastPDF(visual,
                tercile.limits = tas_hcst_terciles[3:1,] - 273.15,
                extreme.limits = tas_hcst_extremes[3:1,] - 273.15,
                var.name = "Air temperature at surface",
                title = "Bias Corrected forecasts at Lappland (4)",
                fcst.names = month.name[rev(as.numeric(months_in_advance))],
                obs = as.vector(tas_obs_areave$data) - 273.15)
  dev.off()
}
print('HERE')


