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

### Creation date: 29th August 2022
# Author: N. Pérez-Zanón
# Refer CSTools package and manuscript when using it.
# Annual verification for Lapland
# ----------------------------------------

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$_$dates$.nc")

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

sdates <- sprintf("%02d", 1:12)
#---------------------
ref_period_ini <- 1993
ref_period_end <- 2016
frpss <- NULL
bss10 <- NULL
bss90 <- NULL
frpss_sig <- NULL
bss10_sig <- NULL
bss90_sig <- NULL
frpss_bc <- NULL
bss10_bc <- NULL
bss90_bc <- NULL
frpss_bc_sig <- NULL
bss10_bc_sig <- NULL
bss90_bc_sig <- NULL

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

# for (sd in 1:length(sdates)) {
 sd <- 12
  # Load the system to verify:
  hcst_sdates <- paste0(ref_period_ini:ref_period_end, sdates[sd], "01")
  tas_hcst <- Start(dat = exp_path,
                    var = 'tas',
                    sdate = hcst_sdates,
                    member = 'all',
                    time = indices(2:4),
                    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'),
                                    member = c('member', 'ensemble')),
                    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)
  dates <- attributes(tas_hcst)$Variable$common$time 

  dim_sdate <- dim(tas_hcst)['sdate']
  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
  # Corresponding array:
  dates <- paste0(format(dates, "%Y"), format(dates, "%m"))
  dim(dates) <- c(sdate = length(ref_period_ini:ref_period_end), time = 3)
  # Load the reference for the hindcaast period:
  tas_ref <- Start(dat = obs_path,
             var = 'tas',
             dates = dates,
             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),
             retrieve = TRUE)
  # Subset Lappland region:
  ys <- attributes(tas_ref)$Variable$common$y
  xs <- attributes(tas_ref)$Variable$common$x

  tas_ref_lapland <- Apply(list(tas_ref, mask_fmi$mask), target_dims = c('x', 'y'),
                              function(x, y) { 
                               mean(x[which(y == 1)], na.rm = TRUE)})$output1
  
  metric <- Apply(list(tas_hcst_lapland, tas_ref_lapland), 
                  target_dims = list(c('sdate', 'member'), c('sdate')),
                  fun = function(x, y) {
                    frpss <- easyVerification::veriApply("FairRpss",
                      fcst = x, obs = y,
                      prob = c(1/3, 2/3), tdim = 1, ensdim = 2)
                    bss10 <- easyVerification::veriApply("FairRpss",
                      fcst = x, obs = y,
                      prob = (1/10), tdim = 1, ensdim = 2)
                    bss90 <- easyVerification::veriApply("FairRpss",
                      fcst = x, obs = y,
                      prob = (9/10), tdim = 1, ensdim = 2)
                   return(list(frpss = frpss$skillscore, 
                               frpss_sd = frpss$skillscore.sd,
                               bss10 = bss10$skillscore, 
                               bss10_sd = bss10$skillscore.sd,
                               bss90 = bss90$skillscore,
                               bss90_sd = bss90$skillscore.sd))}, ncores = 4)
  
  signif_95 <- Apply(metric, target_dims = c('dat'),
                  fun = function(rp, rp.sd, b1, b1.sd, b9, b9.sd, my.pvalue) {
                  rp <- (rp > 0) & (rp > rp.sd * qnorm(1 - my.pvalue))
                  b1 <- (b1 > 0) & (b1 > b1.sd * qnorm(1 - my.pvalue))
                  b9 <- (b9 > 0) & (b9 > b9.sd * qnorm(1 - my.pvalue))
                  return(list(rpss.sig = rp, bss10.sig = b1, bss90.sig = b9))},
                  my.pvalue = 0.05, ncores = 4)
  frpss <- abind::abind(frpss, as.vector(metric$frpss), along = 2)
  names(dim(frpss)) <- c('ftime', 'sdate')
  bss10 <- abind::abind(bss10, as.vector(metric$bss10), along = 2)
  names(dim(bss10)) <- c('ftime', 'sdate')
  bss90 <- abind::abind(bss90, as.vector(metric$bss90), along = 2)
  names(dim(bss90)) <- c('ftime', 'sdate')
  frpss_sig <- abind::abind(frpss_sig, as.vector(signif_95$rpss.sig), along = 2)
  names(dim(frpss_sig)) <- c('ftime', 'sdate')
  bss10_sig <- abind::abind(bss10_sig, as.vector(signif_95$bss10.sig), along = 2)
  names(dim(bss10_sig)) <- c('ftime', 'sdate')
  bss90_sig <- abind::abind(bss90_sig, as.vector(signif_95$bss90.sig), along = 2)
  names(dim(bss90_sig)) <- c('ftime', 'sdate')
  
  #tas_hcst_bc <- BiasCorrection(exp = tas_hcst_lapland,
  #                              obs = tas_ref_lapland)
 # CHECKING Dynamical Bias Correction:
 source("/esarchive/scratch/nperez/git/cstools/R/CST_ProxiesAttractor.R")
 source("/esarchive/scratch/nperez/git/cstools/R/CST_MergeDims.R")
 source("/esarchive/scratch/nperez/git/cstools/R/CST_SplitDim.R")
  library(ClimProjDiags)
  library(abind)
  
  qm = 0.6
  prox <- ProxiesAttractor(tas_hcst, quanti = qm,
                           time_dims = c('time', 'sdate'))
  proxo <- ProxiesAttractor(tas_ref, quanti = qm, time_dims = c('time', 'sdate'),
                            spatial_dims = c('x', 'y'))

  str(prox)
  dim(prox$dim)
  dev.new()
  time = c(1:length(prox$theta))
  layout(matrix(c(1, 3, 2, 3), 2, 2))
  plot(1:24, rep(0, 24), ylim = range(prox$dim), type = 'n', xlab = 'sdate', ylab = 'd',
       main = paste("local dimension", qm))
  for (m in 1:25) {
    matplot(1:24, prox$dim[,1,1,m,], col = 1:3, lty = 1, 
            pch = 16, type = 'b', add = TRUE)
  }
  plot(1:24, rep(0, 24), ylim = range(prox$theta), type = 'n', xlab = 'sdate', ylab = 'theta', 
       main = paste("theta", qm))
   for (m in 1:25) {   
   matplot(1:24, prox$theta[,1,1,m,], type = 'p', pch = 16, col =1:3, add = T)
  }
  plot(1, ylim = range(prox$theta), xlim = range(prox$dim),
       type = 'n', xlab = 'local dimension', ylab = 'theta',
       main = paste("Proxies of the Attractor", qm))
   for (m in 1:25) {   
   matplot(prox$dim[,1,1,m,], prox$theta[,1,1,m,], type = 'p', pch = 16, col =1:3, add = T)
  }

# Plot density:
range(prox$theta)
range(prox$dim)
bins_theta <- seq(floor(range(prox$theta)[1]*10)/10,
                  floor(range(prox$theta)[2]*10)/10, 0.1)
bins_dim <- seq(floor(range(prox$dim)[1]*10)/10, 
                floor(range(prox$dim)[2]*10)/10, 0.1)
density <- array(NA, c(thetas = length(bins_theta), dims = length(bins_dim)))
for (i in 1:(length(bins_theta)-1)) {
  for (j in 1:(length(bins_dim)-1)) {
    pos_dim <- which(prox$dim > bins_dim[j] & prox$dim <= bins_dim[j+1])
    pos_theta <- which(prox$theta > bins_theta[i] & prox$theta <= bins_theta[i+1])
    density[i,j] <- sum(pos_dim %in% pos_theta)
  }
}
library(fields)
dev.new()
image.plot(bins_theta, bins_dim, density[-length(bins_theta), -length(bins_dim)], 
           main = paste("Counts", qm))



  source("/esarchive/scratch/nperez/git/cstools/R/Predictability.R")
  predyn <- Predictability(dim = prox$dim,
                           theta = prox$theta, time_dims = c('time', 'sdate'))
  library(scales)
  dev.new()
  layout(c(1,2))
    plot(prox$dim[,1,1,,1], predyn$dyn_scores[,1,1,,1], col =  alpha(1, 0.3), pch = 16,  
         xlim = range(prox$dim), ylim = range(predyn$dyn_scores), 
          main = paste('Dynamical scores', qm), xlab = 'local dim', ylab = 'dynamical scores')
    points(prox$dim[,1,1,,2], predyn$dyn_scores[,1,1,,2]+0.02, col =  alpha(2, 0.3), pch=16)
    points(prox$dim[,1,1,,3], predyn$dyn_scores[,1,1,,3]+0.04, col = alpha(3, 0.3), pch = 16) 
    plot(prox$theta[,1,1,,1], predyn$dyn_scores[,1,1,,1], col = 1, pch = 1, 
         xlim = range(prox$theta), ylim = range(predyn$dyn_scores),
          main = '', xlab = 'theta', ylab = 'dynamical scores')
    points(prox$theta[,1,1,,2], predyn$dyn_scores[,1,1,,2], col = 2, pch = 1)
    points(prox$theta[,1,1,,3], predyn$dyn_scores[,1,1,,3], col = 3, pch = 1)


  matplot(1:72, predyn$dyn_scores[,1,1,1:5], type = 'b')
  metric <- Apply(list(tas_hcst_bc, tas_ref_lapland),
                  target_dims = list(c('sdate', 'member'), c('sdate')),
                  fun = function(x, y) {
                    frpss <- easyVerification::veriApply("FairRpss",
                      fcst = x, obs = y,
                      prob = c(1/3, 2/3), tdim = 1, ensdim = 2)
                    bss10 <- easyVerification::veriApply("FairRpss",
                      fcst = x, obs = y,
                      prob = (1/10), tdim = 1, ensdim = 2)
                    bss90 <- easyVerification::veriApply("FairRpss",
                      fcst = x, obs = y,
                      prob = (9/10), tdim = 1, ensdim = 2)
                     return(list(frpss = frpss$skillscore,
                               frpss_sd = frpss$skillscore.sd,
                               bss10 = bss10$skillscore,
                               bss10_sd = bss10$skillscore.sd,
                               bss90 = bss90$skillscore,
                               bss90_sd = bss90$skillscore.sd))}, ncores = 4)
  signif_95 <- Apply(metric, target_dims = c('dat'),
                  fun = function(rp, rp.sd, b1, b1.sd, b9, b9.sd, my.pvalue) {
                  rp <- (rp > 0) & (rp > rp.sd * qnorm(1 - my.pvalue))
                  b1 <- (b1 > 0) & (b1 > b1.sd * qnorm(1 - my.pvalue))
                  b9 <- (b9 > 0) & (b9 > b9.sd * qnorm(1 - my.pvalue))
                  return(list(rpss.sig = rp, bss10.sig = b1, bss90.sig = b9))},
                  my.pvalue = 0.05, ncores = 4)

  frpss_bc <- abind::abind(frpss_bc, as.vector(metric$frpss), along = 2)
  names(dim(frpss_bc)) <- c('ftime', 'sdate')
  bss10_bc <- abind::abind(bss10_bc, as.vector(metric$bss10), along = 2)
  names(dim(bss10_bc)) <- c('ftime', 'sdate')
  bss90_bc <- abind::abind(bss90_bc, as.vector(metric$bss90), along = 2)
  names(dim(bss90_bc)) <- c('ftime', 'sdate')
  frpss_bc_sig <- abind::abind(frpss_bc_sig, as.vector(signif_95$rpss.sig), along = 2)
  names(dim(frpss_bc_sig)) <- c('ftime', 'sdate')
  bss10_bc_sig <- abind::abind(bss10_bc_sig, as.vector(signif_95$bss10.sig), along = 2)
  names(dim(bss10_bc_sig)) <- c('ftime', 'sdate')
  bss90_bc_sig <- abind::abind(bss90_bc_sig, as.vector(signif_95$bss90.sig), along = 2)
  names(dim(bss90_bc_sig)) <- c('ftime', 'sdate')

}
# RAW:
t(round(frpss, 3))
t(round(bss10, 3))
t(round(bss90, 3))

t(frpss_sig)
t(bss10_sig)
t(bss90_sig)


# BC:
t(round(frpss_bc, 3))
t(round(bss10_bc, 3))
t(round(bss90_bc, 3))

t(frpss_bc_sig)
t(bss10_bc_sig)
t(bss90_bc_sig)

 


print("Saveing RDS")
saveRDS(list(metric_raw = list(frpss = frpss, frpss.pv = frpss_sig, 
                               bss10 = bss10, bss10.pv = bss10_sig,
                               bss90 = bss90, bss90.pv = bss90_sig),
             metric_bc = list(frpss_bc = frpss_bc, frpss_bc.pv = frpss_bc_sig, 
                              bss10_bc = bss10_bc, bss10_bc.pv = bss10_bc_sig,
                              bss90_bc = bss90_bc, bss90_bc.pv = bss90_bc_sig)),
        file = paste0(output_dir, "FMI_Lapland_Tas_skillscores_Annual_1993-2016.RDS"))


