#!/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.
# ----------------------------------------
# ROVANIEMI: 66.5039° N, 25.7294° E
#             UTM Easting	443474.771302007
#             UTM Northing	7376658.41204131
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 = "FMI",
                 path = "/esarchive/obs/fmi/climgrid/monthly_mean/$var$/$var$_$sdate$$ftime$.nc")

mask_fmi <- readRDS(paste0(working_dir, "masks/FMI_Lapland_mask.RDS"))
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:
tas_ref <- Start(dat = obs_path,
             var = 'tas',
             sdate = as.character(ref_period_ini:ref_period_end),
             ftime = sprintf(target_m),
             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

pos_x_rovaniemi <- max(which(xs < 443474.77130200)) + 1
pos_y_rovaniemi <- min(which(ys < 7376658.41204131))
tas_ref_hcst_rovaniemi <- tas_ref[,,,,pos_y_rovaniemi, pos_x_rovaniemi]
dim_sdate <- dim(tas_ref)['sdate']
dim(tas_ref_hcst_rovaniemi) <- c(sdate = dim_sdate)

  ## check:
 #test <- tas_ref[1,1,1,1,,]
 #test[which(t(mask_fmi$mask) != 1)] <- NA
 #image(xs, sort(ys), t(test))
 #points(xs[pos_x_rovaniemi], sort(ys)[pos_y_rovaniemi], cex = 3)

# ERA5Land Subset region
lat_min <- 65#59
lat_max <- 71
lon_min <- 20
lon_max <- 31

# 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
  
  pos_lat_rovaniemi <- which.min(abs(lats - 66.5039))
  pos_lon_rovaniemi <- which.min(abs(lons - 25.72))
  tas_hcst_rovaniemi <- tas_hcst[,,,,,pos_lat_rovaniemi, pos_lon_rovaniemi]
  

    ## check
    #test <- tas_hcst[1,1,1,1,1,,]
    #test[which(t(mask_era5L$mask) != 1)] <- NA
    #image(lons, lats, t(test))
    #points(lons[pos_lon_rovaniemi], lats[pos_lat_rovaniemi], cex = 3)


  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
  pos_lat_rovaniemi <- which.min(abs(lats - 66.5039))
  pos_lon_rovaniemi <- which.min(abs(lons - 25.72))
  tas_fcst_rovaniemi <- tas_fcst[,,,,,pos_lat_rovaniemi, pos_lon_rovaniemi]



  # SKILL RAW:
  dim(tas_hcst_rovaniemi) <- c(dim_sdate, member = 25)
  dim(tas_ref_hcst_rovaniemi) <- c(dim_sdate)
  dim(tas_fcst_rovaniemi) <- c(sdate = 1, member = 51)

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

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

 
  tas_fcst_bc <- BiasCorrection(exp = tas_hcst_rovaniemi,
                                obs = tas_ref_hcst_rovaniemi, 
                                exp_cor = tas_fcst_rovaniemi)
  tas_fcst_BC[[mm]] <- tas_fcst_bc
  tas_fcst_raw[[mm]] <- tas_fcst_rovaniemi

  # SKILL BC:
  tas_hcst_bc <- BiasCorrection(exp = tas_hcst_rovaniemi,
                                obs = tas_ref_hcst_rovaniemi) 
  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_rovaniemi,
                      prob = c(1/3, 2/3), tdim = 1, ensdim = 2)
  metric[[mm]]$crpss <- easyVerification::veriApply("FairCrpss",
                      fcst = tas_hcst_bc, obs = tas_ref_hcst_rovaniemi,
                      tdim = 1, ensdim = 2)
  metric[[mm]]$bss10 <- easyVerification::veriApply("FairRpss",
                      fcst = tas_hcst_bc, obs = tas_ref_hcst_rovaniemi,
                      prob = (1/10), tdim = 1, ensdim = 2)
  metric[[mm]]$bss90 <- easyVerification::veriApply("FairRpss",
                      fcst = tas_hcst_bc, obs = tas_ref_hcst_rovaniemi,
                      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 = sprintf(target_m),
                 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_rovaniemi <- tas_obs[,,,,pos_y_rovaniemi, pos_x_rovaniemi]

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)))

print("Creating plot")
# Case 4: terciles and extremes are computed with the hcst
#         bias correction is computed on the areave data
#  agg_png(paste0(output_dir, "FMI_Tas_PlotForecast_Rovaniemi_",year, target_m, ".png"),
#    width = 1000, height = 1000, units = 'px',res = 150)
#png(paste0(output_dir, "FMI_Tas_PlotForecast_Rovaniemi_",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 Rovaniemi for", 
                              month.name[as.numeric(target_m)], year),
                fcst.names = subtitles, 
                obs = as.vector(tas_obs_rovaniemi) - 273.15,
                plotfile = paste0(output_dir, "FMI_Rovaniemi_Tas_PlotForecastPDFs",
                                  year, target_m, ".png"))

  #dev.off()

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

print(paste("BC Rovaniemi 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 Rovaniemi 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 Rovaniemi 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 Rovaniemi 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 Rovaniemi 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 Rovaniemi 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 Rovaniemi 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 Rovaniemi 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)))


