#!/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)
library(lubridate)
working_dir <- "/esarchive/scratch/nperez/git/Flor/ArticXchange/"
output_dir <- "/esarchive/scratch/nperez/ArticXChange/round_table/"

setwd(working_dir)

hcst_path <- list(name = "NCEP-CFSv2",
                  path = paste0('/esarchive/exp/ncep/cfs-v2/',
                     'weekly_mean/s2s/$var$_f24h/$var$_$sdate$.nc'))
#fcst_path <- list(name = "ECMWF-S2S",
#                  path = paste0("/esarchive/exp/ecmwf/s2s-monthly_ensfor/",
#                     "weekly_mean/$var$_f24h/$var$_$sdate$.nc"))

obs_path <- list(name = "FMI",
      path = "/esarchive/obs/fmi/climgrid/weekly_mean/$var$_f24h/$var$_$sdate$.nc")

mask_fmi <- readRDS(paste0(working_dir, "masks/FMI_Herding_mask.RDS"))
mask_era5L <- readRDS(paste0(working_dir, "masks/ERA5_Herding_mask.RDS"))

  # Assess forecast from 1 to 3 months in advance
# PARAMETERS TO CHANGE:
year <- 2014
target_m <- '04'
target_w_monday <- '28'
monday_date <- as.Date(paste0(year, target_m, target_w_monday), format = "%Y%m%d")

#---------------------
previous_monday <- monday_date - 7
weeks_in_advance <- previous_monday - c(0, 7, 7*2, 7*3) 
ref_period_ini <- 1999
ref_period_end <- 2018
hcst_period <- ref_period_ini:ref_period_end
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(), list())
metric <- list(list(), list(), list(), list())

# Cross-validation when the year belongs to hindcast period:
if (year %in% ref_period_ini:ref_period_end) {
  hcst_period <- hcst_period[!(year == hcst_period)]
}
  
# ERA5Land Subset region
lat_min <- 64#59
lat_max <- 71
lon_min <- 20
lon_max <- 31
source("/esarchive/scratch/nperez/git/S2S4E-backend-BSC/data-analysis/Vitigeoos/tools.R")
library(abind)
# For each weeks in advance:
for (mm in 1:4) {
  print("Four initializations:")
  print(mm)
  print(paste('Initialization', weeks_in_advance[mm]))
  # Generate the start dates of hindcast period
  print(paste('Hindcast period until', ref_period_end))
  sdates <- get.file_dates(weeks_in_advance[mm])
  hcst_sdates <- apply(sdates, c(1,2), function(x) {paste0(hcst_period, x)})
  names(dim(hcst_sdates))[1] <- 'syear'
  tas_hcst <- Start(dat = hcst_path,
                    var = 'tas',
                    sdate = hcst_sdates,
                    split_multiselected_dims = TRUE,
                    member = 'all', 
                    time = indices(mm),
                    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'),
                    transform_vars = c('lat', 'lon'),
                    synonims = list(member =c('member', 'ensemble'),
                                    lat = c('lat','latitude'),
                                    lon = c('lon','longitude')),
                    return_vars = list(lon = NULL, lat = NULL, 
                                       time = 'sdate'),
                    retrieve = TRUE,
                    num_procs = 2)
  lons <- as.vector(attributes(tas_hcst)$Variable$common$lon)
  lats <- as.vector(attributes(tas_hcst)$Variable$common$lat)                        
  dim(attributes(tas_hcst)$Variable$common$time)
  times_hcst <- attributes(tas_hcst)$Variable$common$time
 
  tas_hcst <- Apply(list(tas_hcst, mask_era5L$mask), 
                         target_dims = c('lon', 'lat'),
                            function(x, y) {
                            mean(x[which(y != 0)], na.rm = TRUE)})$output1

    
  # Load the reference for the hindcaast period:
  #sdates <- get.file_dates(format(weeks_in_advance[mm]+4, "%Y%m%d"))
  #ref_sdates <- apply(sdates, c(1,2), function(x) {paste0(hcst_period, x)})
  #names(dim(ref_sdates))[1] <- 'syear'
  ref_sdates <- format(as.Date(hcst_sdates, "%Y%m%d") + 7*(mm-1), "%Y%m%d")
  dim(ref_sdates) <- dim(hcst_sdates)
  tas_ref <- Start(dat = obs_path,
             var = 'tas',
             sdate = ref_sdates,
             time = indices(1), 
             split_multiselected_dims = TRUE,
             y = 'all',
             x = 'all',
             synonims = list(y = c('y', 'lat', 'latitude'),
                             x = c('x', 'lon', 'longitude')),
             return_vars = list(y = NULL, x = NULL, time = 'sdate'),
             retrieve = TRUE)
  # Subset Lappland region:
  ys <- attributes(tas_ref)$Variable$common$y
  xs <- attributes(tas_ref)$Variable$common$x
  times_ref <- attributes(tas_ref)$Variable$common$time
  # All herding area:
  tas_ref <- Apply(list(tas_ref, mask_fmi$mask),
                           target_dims = c('x', 'y'),
                              function(x, y) {
                               mean(x[which(y != 0)], na.rm = TRUE)})$output1
  # Metrics:
  source("/esarchive/scratch/nperez/cstools/R/CST_BiasCorrection.R")
library(ClimProjDiags)
  tas_hcst_bc <- BiasCorrection(tas_hcst, tas_ref, 
                                sdate_dim = 'syear', window_dim = 'sday',
                                na.rm = TRUE)
  rpss <- mean(RPSS(tas_hcst_bc, tas_ref, time_dim = 'syear', Fair = TRUE)$rpss, na.rm = TRUE) 


  fcst_sdates <- c(paste0(year, format(weeks_in_advance[mm]-4, "%m%d")),
                   paste0(year, format(weeks_in_advance[mm], "%m%d")),
                   paste0(year, format(weeks_in_advance[mm]+3, "%m%d")))
  dim(fcst_sdates) <- c(syear = 1, sday = 3)
  tas_fcst <- Start(dat = hcst_path,
                    var = 'tas',
                    sdate = fcst_sdates,
                    member = 'all',
                    time = indices(mm),
                    split_multiselected_dims = TRUE,
                    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'),
                    transform_vars = c('lat', 'lon'),
                    synonims = list(member = c('member', 'ensemble'),
                                    lat = c('lat','latitude'),
                                    lon = c('lon','longitude')),
                    return_vars = list(lon = NULL, lat = NULL, time = 'sdate'),
                    retrieve = TRUE)
  lons <- as.vector(attributes(tas_fcst)$Variable$common$lon)
  lats <- as.vector(attributes(tas_fcst)$Variable$common$lat)
  times_fcst <- attributes(tas_fcst)$Variable$common$time

  tas_fcst <- Apply(list(tas_fcst, mask_era5L$mask),
                         target_dims = c('lon', 'lat'),
                            function(x, y) {
                            mean(x[which(y != 0)], na.rm = TRUE)})$output1
  tas_fcst_bc <- BiasCorrection(tas_hcst, tas_ref, tas_fcst,
                                sdate_dim = 'syear', window_dim = 'sday',
                                na.rm = TRUE)

  tas_hcst_terciles <- rbind(tas_hcst_terciles,
                           quantile(tas_hcst_bc, c(0.3, 0.6), na.rm = TRUE))
  tas_hcst_extremes <- rbind(tas_hcst_extremes,
                          quantile(tas_hcst_bc, c(0.1, 0.9), na.rm = TRUE))
  tas_fcst_BC[[mm]] <- tas_fcst_bc

  metric[[mm]] <- rpss

}



 visual <- data.frame(week4 = as.vector(tas_fcst_BC[[4]][,1,1,1,5,1]) - 273.15,
                      week3 = as.vector(tas_fcst_BC[[3]][,1,1,1,5,1]) - 273.15,
                      week2 = as.vector(tas_fcst_BC[[2]][,1,1,1,5,1]) - 273.15,
                      week1 = as.vector(tas_fcst_BC[[1]][,1,1,1,5,1]) - 273.15)

tas_obs <- Start(dat = obs_path,
                 var = 'tas',
                 sdate = format(monday_date, "%Y%m%d"),
                 ftime = indices(1),
                 y = 'all',
                 x = 'all',
                 synonims = list(y = c('y', 'lat', 'latitude'),
                                 x = c('x', 'lon', 'longitude'),
                                 ftime = c('ftime', 'time')),
                 return_vars = list(y = NULL, x = NULL),
                 retrieve = TRUE)
# Subset Lappland region:
ys <- attributes(tas_ref)$Variable$common$y
xs <- attributes(tas_ref)$Variable$common$x

tas_obs_herd <- Apply(list(tas_obs, mask_fmi$mask), target_dims = c('x', 'y'),
                              function(x, y) {
                               mean(x[which(y != 0)], na.rm = TRUE)})$output1


subtitles <- c(
   paste0('Lead time: 4 weeks\n Start date: ', 
          weeks_in_advance[4], '-01\n fRPSS:', round(metric[[4]], 3)),
   paste0('Lead time: 3 weeks\n Start date: ', 
          weeks_in_advance[3], '-01\n fRPSS:', round(metric[[3]], 3)),
   paste0('Lead time: 2 weeks\n Start date: ', 
          weeks_in_advance[2], '-01\n fRPSS:', round(metric[[2]], 3)),
   paste0('Lead time: 1 weeks\n Start date: ', 
          weeks_in_advance[1], '-01\n fRPSS:', round(metric[[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, "FMI_Herding_Tas_Subseasonal_PlotForecast_v2_",monday_date, ".png"),
    width = 1900, height = 1400, units = 'px',res = 200)
#png(paste0(output_dir, "FMI_Tas_PlotForecast_Lapland_",year, target_m, ".png"),
#    width = 1000, height = 1000, units = 'px',res = 150)

source("/esarchive/scratch/nperez/git/cstools/R/PlotForecastPDF.R")
library(ggplot2)
library(reshape2)
library(data.table)
library(plyr)
  PlotForecastPDF(visual,
                tercile.limits = tas_hcst_terciles[4:1,] - 273.15,
                extreme.limits = tas_hcst_extremes[4:1,] - 273.15,
                add.ens = 'no',
                color.set = "vitigeoss",
                var.name = "Air temperature at surface",
                title = paste("Seasonal forecasts at Herding region for", 
                              month.name[as.numeric(target_m)], year),
                fcst.names = subtitles, #month.name[rev(as.numeric(months_in_advance))],
                obs = as.vector(tas_obs_herd) - 273.15)#,
                #plotfile = paste0(output_dir, "FMI_Herding_Tas_Subseasonal_PlotForecast_",
                #                  year, target_m, target_w_monday, ".png"))

  dev.off()

print("Saveing RDS")

print(paste("BC Herding FMI", 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 Herding FMI", 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 FMI", 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 FMI", 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 FMI", 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 FMI", 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 FMI", 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 FMI", 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)))

