#!/usr/bin/env Rscript

rm(list=ls()); gc();

### Creation date: 3rd October 2022
# 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(ggplot2)
library(RColorBrewer)

working_dir <- "/esarchive/scratch/nperez/git/Flor/ArticXchange/"
output_dir <- "/esarchive/scratch/nperez/ArticXChange/"

setwd(working_dir)
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_Herding_mask.RDS"))
mask_era5L <- readRDS(paste0(working_dir, "masks/ERA5_Herding_mask.RDS"))
polys <- unique(as.vector(mask_fmi$mask))[-1]
dim(polys) <- c(poly = length(polys))


  # 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

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

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

mask.file_4 <- paste0('region_exploration/LAU/LAU_RG_01M_2020_4326.shp')
shp_4 <- rgdal::readOGR(mask.file_4)
shp_4 <- subset(shp_4, CNTR_CODE == "FI")
shp_herding <- subset(shp_4, LAU_ID %in% sprintf("%03d", polys))
gusa <- map_data("world")

plot_list <- list()
# 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_herd <- Apply(list(tas_hcst, mask_era5L$mask, polys), 
                         target_dims = list(c('lon', 'lat'), c('lon', 'lat'), NULL),
                            function(x, y, z) {
                            mean(x[which(y == z)], 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_herd <- Apply(list(tas_fcst, mask_era5L$mask, polys), 
                         target_dims = list(c('lon', 'lat'), c('lon', 'lat'), NULL),
                            function(x, y, z) {
                            mean(x[which(y == z)], 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:
  tas_hcst_herd <- drop(tas_hcst_herd)
  tas_ref_hcst_herd <- drop(tas_ref_hcst_herd)
  dim(tas_fcst_herd) <- c(sdate = 1, member = 51, poly = 31)
  names(dim(tas_hcst_herd))[2] <- 'member' 

  tas_hcst_herd <- Reorder(tas_hcst_herd, c('poly', 'sdate', 'member'))
  tas_ref_hcst_herd <- Reorder(tas_ref_hcst_herd, c('poly', 'sdate'))
  tas_fcst_herd <- Reorder(tas_fcst_herd, c('poly', 'sdate', 'member'))

  terciles <- Apply(list(tas_hcst_herd), target_dims = c('sdate', 'member'),
                    function(x) {quantile(as.vector(x), c(0.33,0.66))})$output1
  tas_hcst_terciles <- abind::abind(tas_hcst_terciles, terciles, along = 3)
  extremes <- Apply(list(tas_hcst_herd), target_dims = c('sdate', 'member'),
                    function(x) {quantile(as.vector(x), c(0.1,0.9))})$output1
  tas_hcst_extremes <- abind::abind(tas_hcst_extremes, extremes, along = 3)

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

 
  tas_fcst_bc <- BiasCorrection(exp = tas_hcst_herd,
                                obs = tas_ref_hcst_herd, 
                                exp_cor = tas_fcst_herd)
  tas_fcst_BC[[mm]] <- tas_fcst_bc
  tas_fcst_raw[[mm]] <- tas_fcst_herd

  # SKILL BC:
  tas_hcst_bc <- BiasCorrection(exp = tas_hcst_herd,
                                obs = tas_ref_hcst_herd) 
  tas_hcst_bc <- Reorder(tas_hcst_bc, c('poly', 'sdate', 'member'))
 
  terciles_bc <- Apply(list(tas_hcst_bc), target_dims = c('sdate', 'member'),
                    function(x) {quantile(as.vector(x), c(0.33,0.66))},
                    output_dims = 'prob')$output1
  tas_hcst_bc_terciles <- abind::abind(tas_hcst_bc_terciles, terciles_bc, along = 3)
  extremes_bc <- Apply(list(tas_hcst_bc), target_dims = c('sdate', 'member'),
                    function(x) {quantile(as.vector(x), c(0.1,0.9))})$output1
  tas_hcst_bc_extremes <- abind::abind(tas_hcst_bc_extremes, extremes_bc, along = 3)

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

# Plot Mostlikely Map

# Data.frame to display positive skill
## Function match is used to reorder the poly dimension accordingly to the shp_file order
  FRPSS <- metric[[mm]]$frpss$skillscore
  FRPSS <- FRPSS[match(shp_herding@data$LAU_ID, sprintf("%03d", polys))]
  
  df.frpss <- data.frame(skill = FRPSS[which(FRPSS < 0)], 
                         LAU_ID = shp_herding@data$LAU_ID[which(FRPSS < 0)])
  shp_nosk <- subset(shp_herding, LAU_ID %in% shp_herding@data$LAU_ID[which(FRPSS < 0)]) 
  shp_frpss <- merge(shp_nosk, df.frpss, by.x="LAU_ID",by.y="LAU_ID")
# Magic lines below that converts SpatialPolygonsDataFrame to normal data frame
  shp_frpss@data$id <- rownames(shp_frpss@data)
  df.points <- fortify(shp_frpss, region = "id")
  df.skill <- plyr::join(df.points, shp_frpss@data, by = "id")
  
# Probabilities of each tercile:
  mostL <- Apply(list(tas_fcst_bc, terciles_bc), target_dims = list('member', 'prob'),
                fun = function(x, y) { 
                    prob <- c(sum(x<y[1]), sum(x>= y[1] & x <= y[2]), sum(x>y[2]))
                    prob <- prob/length(x) * 100
                    }, output_dims = 'bin')$output1
  mostL <- mostL[,1, match(shp_herding@data$LAU_ID, sprintf("%03d", polys))]

## Data frames for each category  
  pos <- apply(mostL, 2, which.max)
  df.bn <- data.frame(ter = mostL[1, which(pos == 1)], 
                      LAU_ID = shp_herding@data$LAU_ID[pos == 1])
  df.nor <- data.frame(ter = mostL[2, which(pos == 2)],
                      LAU_ID = shp_herding@data$LAU_ID[pos == 2])
  df.an <- data.frame(ter = mostL[3, which(pos == 3)],
                      LAU_ID = shp_herding@data$LAU_ID[pos == 3])

  shp_df1 <- subset(shp_herding, LAU_ID %in% shp_herding@data$LAU_ID[pos == 1])
  shp_df2 <- subset(shp_herding, LAU_ID %in% shp_herding@data$LAU_ID[pos == 2])
  shp_df3 <- subset(shp_herding, LAU_ID %in% shp_herding@data$LAU_ID[pos == 3])

# In case a category is empty:
df.df1 <- df.df2 <- df.df3 <- NULL
if (any(pos == 1)) { 
shp_df1 <- merge(shp_df1, df.bn, by.x="LAU_ID", by.y="LAU_ID")
# Magic lines below that converts SpatialPolygonsDataFrame to normal data frame
  shp_df1@data$id <- rownames(shp_df1@data)
  df.points1 <- fortify(shp_df1, region = "id")
  df.df1 <- plyr::join(df.points1, shp_df1@data, by = "id")
} 
if (any(pos == 2)) {
shp_df2 <- merge(shp_df2, df.nor, by.x="LAU_ID", by.y="LAU_ID")
# Magic lines below that converts SpatialPolygonsDataFrame to normal data frame
  shp_df2@data$id <- rownames(shp_df2@data)
  df.points2 <- fortify(shp_df2, region = "id")
  df.df2 <- plyr::join(df.points2, shp_df2@data, by = "id")
} 
if (any(pos == 3)) {
shp_df3 <- merge(shp_df3, df.an, by.x="LAU_ID", by.y="LAU_ID")
# Magic lines below that converts SpatialPolygonsDataFrame to normal data frame
  shp_df3@data$id <- rownames(shp_df3@data)
  df.points3 <- fortify(shp_df3, region = "id")
  df.df3 <- plyr::join(df.points3, shp_df3@data, by = "id")
}
if (is.null(df.df1)) {
  if (is.null(df.df2)) {
       df.df1 <- df.df2 <- df.df3
       df.df1$ter <- df.df2$ter <- 0.1 
  } else {
       df.df1 <- df.df2
       df.df1$ter <- 0.1
  }
}
if (is.null(df.df2)) {
  if (is.null(df.df3)) {
       df.df2 <- df.df1
       df.df2$ter <- 0.1
  } else {
     df.df2 <- df.df1
     df.df2$ter <- 0.1
  }
}
if (is.null(df.df3)) {
  if (is.null(df.df1)) {
    df.df3 <- df.df2
    df.df3$ter <- 0.1
  } else {
    df.df3 <- df.df1
    df.df3$ter <- 0.1
  }
}
 


library(ggnewscale)

p <-  ggplot() +
  geom_polygon(aes(long, lat, group = group),fill = "lightgrey", data = gusa) +
        coord_map(projection = "stereographic",
                  xlim = c(lon_min, lon_max),
                  ylim = c(lat_min, lat_max)) 
p <- p +
  geom_polygon(data = df.df1, aes(x = long, y = lat, fill = cut(ter,
                         breaks = seq(40, 100, 10)),
             group = group), color = "grey", size = 0.08) +
       scale_fill_manual(values = brewer.pal(7, 'Blues'),
                          drop = FALSE, name = 'Below normal', na.translate = F,                          guide = guide_legend(order = 3))
p <- p +
  new_scale("fill") +
  geom_polygon(data = df.df2, aes(x = long, y =lat, fill = cut(ter,
                         breaks = seq(40, 100, 10)),
            group = group), na.rm = T, color = "grey",size = 0.08) +
       scale_fill_manual(values = brewer.pal(7, 'Greens'),
                          drop = FALSE, name = 'Normal',na.translate = F,
                          guide = guide_legend(order = 2)) 
p <- p +
  new_scale("fill") +
  geom_polygon(data = df.df3, aes(x = long, y =lat, fill = cut(ter,
                         breaks = seq(40, 100, 10)),
            group = group), na.rm = T, color = "grey",size = 0.08) +
       scale_fill_manual(values = brewer.pal(7, 'Reds'),
                          drop = FALSE, name = 'Above normal',na.translate = F,
                          guide = guide_legend(order = 1))
p <- p +
  new_scale("fill") +
  geom_polygon(data = df.skill, aes(x=long, y = lat, fill = "darkgrey", group = group)) +
  scale_fill_manual(values = "darkgrey", label = 'No skill', name = '')

p <- p +
    geom_polygon(aes(x = long, y = lat, group = group), fill = NA,
                    color = "grey", size = 0.5, data = shp_herding) +
   theme_bw() +
        theme(panel.background = element_rect(fill = 'azure'),
              text = element_text(family = "Times")) +  # Colour of ocean
        xlab('Longitude') + ylab('Latitude') +
        scale_x_continuous(breaks = seq(-12, 45, 4),labels = waiver()) +
        scale_y_continuous(breaks = seq(32, 70, 4),labels = waiver()) +
        ggtitle(paste(mm, "months in advance"))# +

plot_list[[mm]] <- p

}
library(ggpubr)

plot <- ggarrange(plotlist = rev(plot_list), ncol = 3, nrow = 1, common.legend = TRUE, legend = "bottom")
plot <- annotate_figure(plot,
        top = text_grob(paste("SEAS5 / tas / Most likely terciles",
                        month.name[as.numeric(target_m)], year, "/ FMI",
                        ref_period_ini, "-", ref_period_end),
                        color = "black", face = "bold", size = 14))
ggsave(paste0("SEAS5_FMI_Herding_MostLikely_", target_m, "_", year, ".jpg"), 
       plot, width = 9, height = 6.5, bg = 'white', dpi = 300)
