#!/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.
# ----------------------------------------

library(CSTools)
library(s2dv)
library(ragg)
library(multiApply)
library(startR)
library(sf)
library(rgeos)
library(multiApply)
working_dir <- "/esarchive/scratch/nperez/git/Flor/ArticXchange/"
output_dir <- "/esarchive/scratch/nperez/ArticXChange/"

exp_path <- list(name = "ECMWFS5",
                 path = "/esarchive/exp/ecmwf/system5c3s/monthly_mean/$var$_f6h/$var$_$sdate$.nc")
obs_path <- list(name = "ERA5Land",
                 path = "/esarchive/recon/ecmwf/era5land/monthly_mean/$var$_f1h/$var$_$sdate$$ftime$.nc")

mask_era5L <- readRDS(paste0(working_dir, "masks/ERA5_Lapland_mask.RDS"))

  # Assess forecast from 1 to 3 months in advance
# PARAMETERS TO CHANGE:
year <- 2020
target_m <- '06'
#---------------------
months_in_advance <- sprintf("%02d", (as.numeric(target_m)-1):(as.numeric(target_m)-3))
ref_period_ini <- 1993
ref_period_end <- 2016
tas_fcst_BC <- list()
tas_fcst_raw <- list()
tas_hcst_terciles <- NULL
tas_hcst_extremes <- NULL
tas_hcst_bc_terciles <- NULL
tas_hcst_bc_extremes <- NULL
metric_raw <- list(list(), list(), list())
metric <- list(list(), list(), list())

# Load the reference for the hindcaast period:
# ERA5Land Subset region
lat_min <- 65#59
lat_max <- 71
lon_min <- 20
lon_max <- 31

tas_ref <- Start(dat = obs_path,
             var = 'tas',
             sdate = as.character(ref_period_ini:ref_period_end),
             ftime = target_m,
             lat = values(list(lat_min, lat_max)),
             lat_reorder = Sort(),
             lon = values(list(lon_min, lon_max)),
             lon_reorder = CircularSort(-180, 180),
             synonims = list(lat=c('lat','latitude'),
                             lon=c('lon','longitude')),
             return_vars = list(lon = NULL, lat = NULL),
             retrieve = TRUE)
# Subset Lappland region:
lats <- attributes(tas_ref)$Variable$common$lat
lons <- attributes(tas_ref)$Variable$common$lon

tas_ref_hcst_lapland <- Apply(list(tas_ref, mask_era5L$mask),
                              target_dims = c('lon', 'lat'),
                              function(x, y) { 
                                mean(x[which(y == 1)], na.rm = TRUE)})$output1
dim_sdate <- dim(tas_ref)['sdate']
  ## check:
  #test <- tas_ref[1,1,1,1,,]
  #s2dv::PlotEquiMap(test, lon = lons, lat = lats, filled.c = F)
  #test_masked <- test
  #test_masked[which(t(mask_era5L$mask) != 1)] <- NA
  #dev.new()
  #s2dv::PlotEquiMap(test_masked, lon = lons, lat = lats, filled.c = F)
  #mean(test_masked, na.rm = T)  
  #tas_ref_hcst_lapland[1,1,1,1]


# 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', ref_period_end))
  hcst_sdates <- paste0(ref_period_ini:ref_period_end, months_in_advance[mm], '01') 
  tas_hcst <- Start(dat = exp_path,
                    var = 'tas',
                    sdate = hcst_sdates,
                    ensemble = 'all', 
                    time = indices(mm + 1),
                    lat = values(list(lat_min, lat_max)),
                    lat_reorder = Sort(),
                    lon = values(list(lon_min, lon_max)),
                    lon_reorder = CircularSort(-180, 180),
                    transform = CDORemapper,
                    transform_extra_cells = 0,
                    transform_params = list(grid =
  '/esarchive/recon/ecmwf/era5land/monthly_mean/tas_f1h/tas_201401.nc',
                                     method = 'conservative',
                                     crop = c(lon_min, lon_max, lat_min, lat_max)),
                    transform_vars = c('lat', 'lon'),
                    synonims = list(lat = c('lat','latitude'),
                                    lon = c('lon','longitude')),
                    return_vars = list(lon = NULL, lat = NULL, time = 'sdate'),
                    retrieve = TRUE)
  lons <- as.vector(attributes(tas_hcst)$Variable$common$lon)
  lats <- as.vector(attributes(tas_hcst)$Variable$common$lat)                        
  attributes(tas_hcst)$Variable$common$time
  
  tas_hcst_lapland <- Apply(list(tas_hcst, mask_era5L$mask), target_dims = c('lon', 'lat'),
                            function(x, y) {
                            mean(x[which(y == 1)], na.rm = TRUE)})$output1
    ## check
    #test <- tas_hcst[1,1,1,1,1,,]
    #s2dv::PlotEquiMap(test, lon = lons, lat = lats, filled.c = F)
    #test_masked <- test
    #test_masked[which(t(mask_era5L$mask) != 1)] <- NA
    #dev.new()
    #s2dv::PlotEquiMap(test_masked, lon = lons, lat = lats, filled.c = F)
    #mean(test_masked, na.rm = T)  
    #tas_hcst_lapland[1,1,1,1,1]


  fcst_sdates <- paste0(year, months_in_advance[mm], '01')
  tas_fcst <- Start(dat = exp_path,
                    var = 'tas',
                    sdate = fcst_sdates,
                    ensemble = 'all',
                    time = indices(mm + 1),
                    lat = values(list(lat_min, lat_max)),
                    lat_reorder = Sort(),
                    lon = values(list(lon_min, lon_max)),
                    lon_reorder = CircularSort(-180, 180),
                    transform = CDORemapper,
                    transform_extra_cells = 0,
                    transform_params = list(grid =
  '/esarchive/recon/ecmwf/era5land/monthly_mean/tas_f1h/tas_201401.nc',
                                     method = 'conservative',
                                     crop = c(lon_min, lon_max, lat_min, lat_max)),
                    transform_vars = c('lat', 'lon'),
                    synonims = list(lat = c('lat','latitude'),
                                    lon = c('lon','longitude')),
                    return_vars = list(lon = NULL, lat = NULL, time = NULL),
                    retrieve = TRUE)
  lons <- as.vector(attributes(tas_fcst)$Variable$common$lon)
  lats <- as.vector(attributes(tas_fcst)$Variable$common$lat)
  attributes(tas_fcst)$Variable$common$time

  tas_fcst_lapland <- Apply(list(tas_fcst, mask_era5L$mask), target_dims = c('lon', 'lat'),
                            function(x, y) {
                            mean(x[which(y == 1)], na.rm = TRUE)})$output1
    ## check
    #test <- tas_fcst[1,1,1,1,1,,]
    #s2dv::PlotEquiMap(test, lon = lons, lat = lats, filled.c = F)
    #test_masked <- test
    #test_masked[which(t(mask_era5L$mask) != 1)] <- NA
    #dev.new()
    #s2dv::PlotEquiMap(test_masked, lon = lons, lat = lats, filled.c = F)
    #mean(test_masked, na.rm = T)  
    #tas_fcst_lapland[1,1,1,1,1]


  # SKILL RAW:
  dim(tas_hcst_lapland) <- c(dim_sdate, member = 25)
  dim(tas_ref_hcst_lapland) <- c(dim_sdate)
  dim(tas_fcst_lapland) <- c(sdate = 1, member = 51)

  tas_hcst_terciles <- rbind(tas_hcst_terciles,
                           quantile(tas_hcst_lapland, c(0.3, 0.6)))
  tas_hcst_extremes <- rbind(tas_hcst_extremes,
                          quantile(tas_hcst_lapland, c(0.1, 0.9)))

  metric_raw[[mm]]$frpss <- easyVerification::veriApply("FairRpss",
                      fcst = tas_hcst_lapland, obs = tas_ref_hcst_lapland,
                      prob = c(1/3, 2/3), tdim = 1, ensdim = 2)
  metric_raw[[mm]]$crpss <- easyVerification::veriApply("FairCrpss",
                      fcst = tas_hcst_lapland, obs = tas_ref_hcst_lapland,
                      tdim = 1, ensdim = 2)
  metric_raw[[mm]]$bss10 <- easyVerification::veriApply("FairRpss",
                      fcst = tas_hcst_lapland, obs = tas_ref_hcst_lapland,
                      prob = (1/10), tdim = 1, ensdim = 2)
  metric_raw[[mm]]$bss90 <- easyVerification::veriApply("FairRpss",
                      fcst = tas_hcst_lapland, obs = tas_ref_hcst_lapland,
                      prob = (9/10), tdim = 1, ensdim = 2)

 
  tas_fcst_bc <- BiasCorrection(exp = tas_hcst_lapland,
                                obs = tas_ref_hcst_lapland, 
                                exp_cor = tas_fcst_lapland)
  tas_fcst_BC[[mm]] <- tas_fcst_bc
  tas_fcst_raw[[mm]] <- tas_fcst_lapland

  # SKILL BC:
  tas_hcst_bc <- BiasCorrection(exp = tas_hcst_lapland,
                                obs = tas_ref_hcst_lapland) 
  tas_hcst_bc <- t(tas_hcst_bc)
  dim(tas_hcst_bc) <- c(dim_sdate, member = 25)

  tas_hcst_bc_terciles <- rbind(tas_hcst_bc_terciles,
                           quantile(tas_hcst_bc, c(0.3, 0.6)))
  tas_hcst_bc_extremes <- rbind(tas_hcst_bc_extremes,
                          quantile(tas_hcst_bc, c(0.1, 0.9)))

  metric[[mm]]$frpss <- easyVerification::veriApply("FairRpss",
                      fcst = tas_hcst_bc, obs = tas_ref_hcst_lapland,
                      prob = c(1/3, 2/3), tdim = 1, ensdim = 2)
  metric[[mm]]$crpss <- easyVerification::veriApply("FairCrpss",
                      fcst = tas_hcst_bc, obs = tas_ref_hcst_lapland,
                      tdim = 1, ensdim = 2)
  metric[[mm]]$bss10 <- easyVerification::veriApply("FairRpss",
                      fcst = tas_hcst_bc, obs = tas_ref_hcst_lapland,
                      prob = (1/10), tdim = 1, ensdim = 2)
  metric[[mm]]$bss90 <- easyVerification::veriApply("FairRpss",
                      fcst = tas_hcst_bc, obs = tas_ref_hcst_lapland,
                      prob = (9/10), tdim = 1, ensdim = 2)

}

 visual <- data.frame(dec = as.vector(tas_fcst_BC[[3]]) - 273.15,
                       jan = as.vector(tas_fcst_BC[[2]]) - 273.15,
                       feb = as.vector(tas_fcst_BC[[1]]) - 273.15)

tas_obs <- Start(dat = obs_path,
                 var = 'tas',
                 sdate = as.character(year),
                 ftime = target_m,
                 lat = values(list(lat_min, lat_max)),
                 lat_reorder = Sort(),
                 lon = values(list(lon_min, lon_max)),
                 lon_reorder = CircularSort(-180, 180),
                 synonims = list(lat = c('lat','latitude'),
                                 lon = c('lon','longitude')),
                 return_vars = list(lon = NULL, lat = NULL),
                 retrieve = TRUE)
# Subset Lappland region:
lons <- attributes(tas_ref)$Variable$common$lat
lats <- attributes(tas_ref)$Variable$common$lon

tas_obs_lapland <- Apply(list(tas_obs, mask_era5L$mask),
                         target_dims = c('lon', 'lat'),
                              function(x, y) {
                               mean(x[which(y == 1)], na.rm = TRUE)})$output1

subtitles <- c(
   paste0('Lad time: 3 months\n Start date: ', year, '-',
          months_in_advance[3], '-01\n fRPSS:', round(metric[[3]]$frpss[[1]], 3)),
   paste0('Lead time: 2 months\n Start date: ', year, '-',
          months_in_advance[2], '-01\n fRPSS:', round(metric[[2]]$frpss[[1]], 3)),
   paste0('Lead time: 1 months\n Start date: ', year, '-', 
          months_in_advance[1], '-01\n fRPSS:', round(metric[[1]]$frpss[[1]], 3)))

# Case 4: terciles and extremes are computed with the hcst
#         bias correction is computed on the areave data
#  agg_png(paste0(output_dir, "ERA5L_Tas_PlotForecast_Lapland_",year, target_m, ".png"),
#    width = 1000, height = 1000, units = 'px',res = 150)
#png(paste0(output_dir, "ERA5L_Tas_PlotForecast_Lapland_",year, target_m, ".png"),
#    width = 1000, height = 1000, units = 'px',res = 150)
  CSTools::PlotForecastPDF(visual,
                tercile.limits = tas_hcst_bc_terciles[3:1,] - 273.15,
                extreme.limits = tas_hcst_bc_extremes[3:1,] - 273.15,
                var.name = "Air temperature at surface",
                title = paste("Seasonal forecasts at Lappland for", 
                              month.name[as.numeric(target_m)], year),
                fcst.names = subtitles, #month.name[rev(as.numeric(months_in_advance))],
                obs = as.vector(tas_obs_lapland) - 273.15,
                plotfile = paste0(output_dir, "ERA5L_Lapland_Tas_PlotForecast_",
                                  year, target_m, ".png"))

 # dev.off()

print("Saveing RDS")
saveRDS(list(metric_raw = metric_raw, metric_bc = metric),
        file = paste0(output_dir, "ERA5L_Lapland_Tas_skillscores_", year, target_m, ".RDS"))

print(paste("BC Lapland ERA5L", year, target_m, "fRPSS (3,2,1 months):"))
print(c(round(metric[[3]]$frpss[[1]], 3),
        round(metric[[2]]$frpss[[1]], 3),
        round(metric[[1]]$frpss[[1]], 3)))

print(paste("Raw Lapland ERA5L", year, target_m, "fRPSS (3,2,1 months):"))
print(c(round(metric_raw[[3]]$frpss[[1]], 3),
        round(metric_raw[[2]]$frpss[[1]], 3),
        round(metric_raw[[1]]$frpss[[1]], 3)))

print(paste("BC Lapland ERA5L", year, target_m, "CRPSS (3,2,1 months):"))
print(c(round(metric[[3]]$crpss[[1]], 3),
        round(metric[[2]]$crpss[[1]], 3),
        round(metric[[1]]$crpss[[1]], 3)))

print(paste("Raw Lapland ERA5L", year, target_m, "CRPSS (3,2,1 months):"))
print(c(round(metric_raw[[3]]$crpss[[1]], 3),
        round(metric_raw[[2]]$crpss[[1]], 3),
        round(metric_raw[[1]]$crpss[[1]], 3)))

print(paste("BC Lapland ERA5L", year, target_m, "BSS10 (3,2,1 months):"))
print(c(round(metric[[3]]$bss10[[1]], 3),
        round(metric[[2]]$bss10[[1]], 3),
        round(metric[[1]]$bss10[[1]], 3)))

print(paste("Raw Lapland ERA5L", year, target_m, "BSS10 (3,2,1 months):"))
print(c(round(metric_raw[[3]]$bss10[[1]], 3),
        round(metric_raw[[2]]$bss10[[1]], 3),
        round(metric_raw[[1]]$bss10[[1]], 3)))

print(paste("BC Lapland ERA5L", year, target_m, "BSS90 (3,2,1 months):"))
print(c(round(metric[[3]]$bss90[[1]], 3),
        round(metric[[2]]$bss90[[1]], 3),
        round(metric[[1]]$bss90[[1]], 3)))

print(paste("Raw Lapland ERA5L", year, target_m, "BSS90 (3,2,1 months):"))
print(c(round(metric_raw[[3]]$bss90[[1]], 3),
        round(metric_raw[[2]]$bss90[[1]], 3),
        round(metric_raw[[1]]$bss90[[1]], 3)))

# EXPLORING RELIABILITY
# source('/home/Earth/vtorralb/vtesis_scripts/reliability/PlotRD.R')
# source('/home/Earth/vtorralb/vtesis_scripts/reliability/ReliabilityDiagramHist.R')
# nmemb <- dim(tas_hcst_lapland)['member']
# #Select Lapland points
# exp <- tas_hcst[1,1,,,1,,]
# exp <- Apply(list(exp, mask_era5L$mask), target_dims = c('lat', 'lon'), 
#              fun = function(x,y) {x[which(y != 1)] <- NA
#                                    return(x)})$output1
# exp <- MergeDims(MergeDims(exp, c('lat', 'lon'), rename = 'gp'), 
#                  c('gp', 'sdate'), rename = 'sdate')
# obs <- Apply(list(tas_ref, mask_era5L$mask), target_dims = c('lat', 'lon'), 
#              fun = function(x,y) {x[which(y != 1)] <- NA
#                                    return(x)})$output1
# obs <- MergeDims(MergeDims(obs, c('lat', 'lon'), rename = 'gp'), 
#                  c('gp', 'sdate'), rename = 'sdate')
# exp <- exp[complete.cases(exp),]
# obs <- obs[complete.cases(obs),]
# dim(obs) <- c(sdate = length(obs), member = 1)
# probs_exp <- convert2prob(exp, prob = 1:2/3)/nmemb
# probs_obs <- convert2prob(obs, prob = 1:2/3)
# rd <- list()
# for(j in 1:3){
#  rd[[j]] <- ReliabilityDiagramHist(probs_exp[,j],
#                                    probs_obs[,j], plot = F, nboot = 500, bins = 10)
#}
# library(abind)
# PlotRD(rd, cons.bars = T, toptitle = 'Reliability Diagrams')

