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

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

mask_fmi <- readRDS(paste0(working_dir, "masks/FMI_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
months <- c('05', '06', '07')
ref_period_ini <- 1993
ref_period_end <- 2021

## DONE: work with months of 30 days in last position:
x <- seq(ymd(paste0(year,"-", months[1], "-01")), 
         ymd(paste0(year, "-", months[length(months)], "-", 
                    days_in_month(as.numeric(months[3])))), by = "1 day")
first_monday <- x[wday(x,label = TRUE) == "Mon" & day(x) <= 7][1] 
index_first_monday <- which(x == first_monday)
last_sunday <- x[wday(x, label = TRUE) == "Sun" & day(x) >= 24]
last_sunday <- last_sunday[length(last_sunday)]
index_last_sunday <- which(x == last_sunday)

# Load the reference for the hindcaast period:
tas_ref <- Start(dat = obs_path,
             var = 'tas',
             sdate = as.character(ref_period_ini:ref_period_end),
             month = sprintf(months),
             time = 'all',
             y = 'all',
             x = 'all',
             time_across = 'month',
             merge_across_dims = TRUE,
             merge_across_dims_narm = TRUE,
             largest_dims_length = TRUE,
             synonims = list(y = c('y', 'lat', 'latitude'),
                             x = c('x', 'lon', 'longitude')),
             return_vars = list(y = NULL, x = NULL),
             retrieve = TRUE)

ys <- attributes(tas_ref)$Variable$common$y
xs <- attributes(tas_ref)$Variable$common$x

#tas2 <- tas_ref[1,1,1,1,,]
#mean(tas2[which(t(mask_fmi$mask) != 0)])
#tas2[which(t(mask_fmi$mask) == 0)] <- NA
#s2dv::PlotEquiMap(tas2, lon = xs, lat = ys, filled.c = F)

# For Herding region:
tas_polys <- Apply(list(tas_ref, mask_fmi$mask),
                   target_dims = list(c('x', 'y'), c('x', 'y')),
                   function(x, y) {
                     mean(x[which(y != 0)], na.rm = TRUE)},
                   ncores = 6)$output1
# By polygons
#tas_polys <- 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)},
#                   ncores = 6)$output1

tas_polys <- tas_polys - 273.15


# selection by weeks:
tas_polys <- Subset(tas_polys, along = 'time', 
                    indices = index_first_monday:index_last_sunday)

source("https://earth.bsc.es/gitlab/external/cstools/-/raw/develop-PlotTimeseriesClim/R/PlotTimeseriesClim.R"))


PlotTimeseriesClim(tas_polys, first_date = as.character(first_monday),
                   ref_period_ini = ref_period_ini,
                   ref_period_end = ref_period_end,
                               time_dim = 'time', sdate_dim = 'sdate',
                               title = "Observed weekly means and climatology",
                               palette = "Blues")

weekly_aggre <- SplitDim(tas_polys, split_dim = 'time', 
                          indices = sort(rep(1:(length(index_first_monday:index_last_sunday)/7), 7)),
                          new_dim_name = 'week')
weekly_means <- MeanDims(weekly_aggre, 'time')

weekly_clim <- MeanDims(weekly_means, 'sdate')
weekly_p10 <- Apply(weekly_means, target_dims = 'sdate',
                    fun = function(x) {quantile(x, 0.10)})$output1
weekly_p90 <- Apply(weekly_means, target_dims = 'sdate',
                    fun = function(x) {quantile(x, 0.90)})$output1
weekly_p33 <- Apply(weekly_means, target_dims = 'sdate',
                    fun = function(x) {quantile(x, 0.33)})$output1
weekly_p66 <- Apply(weekly_means, target_dims = 'sdate',
                    fun = function(x) {quantile(x, 0.66)})$output1

weekly_data <- data.frame(clim = as.vector(weekly_clim),
                          p10 = as.vector(weekly_p10),
                          p90 = as.vector(weekly_p90), 
                          p33 = as.vector(weekly_p33),
                          p66 = as.vector(weekly_p66),
                          week = 1:(length(index_first_monday:index_last_sunday)/7))

daily <- tas_polys[1,1,which(ref_period_ini:ref_period_end == year),]
daily_data <- data.frame(day.x = seq(first_monday, last_sunday, by = "1 day"),
                         data = daily, 
                         week = sort(rep(1:(length(index_first_monday:index_last_sunday)/7), 7)))
week_mean <- aggregate(data~week, data = daily_data, mean)
weekly_data <- cbind(weekly_data, week_mean$data)
colnames(weekly_data)[7] <- 'week_mean'
all <- merge(weekly_data, daily_data, by = 'week')


library(ggplot2)
library(scales)
library(RColorBrewer)
    cols <- colorRampPalette(brewer.pal(9, "Blues"))(6)
p <- ggplot(all, aes(x = day.x)) +
  geom_ribbon(aes(ymin = p10, ymax = p90, group = week, fill = "p10-p90"),
        alpha = 0.7) +  # extremes clim
  geom_ribbon(aes(ymin = p33, ymax = p66, group = week, fill = "p33-p66"),
        alpha = 0.7) +  # terciles clim
  geom_line(aes(y = clim, group = week, color = "climatological mean",
            linetype = "climatological mean"),
          alpha = 1.0, size = 0.7) + 	 # mean clim
  geom_line(aes(y = data, color = "observed daily mean",
                linetype = "observed daily mean"),
            alpha = 1, size = 0.2) + # daily evolution
  geom_line(aes(y = week_mean, group = week, color = "observed weekly mean",
            linetype = "observed weekly mean"),
          alpha = 1, size = 0.7) + 		 # weekly evolution
  theme_bw() + ylab(paste0('tas', " (", "deg.C", ")")) + xlab(NULL) + 
  ggtitle(paste("Herding region", "/", sub(".*/", "", "FMI"), "/",
                year, months[1],#substr(first_month,1,7),
          #substr(last_month, "/", " Ref.: ",
          ref_period_ini,"-",ref_period_end, sep=" ") ) +
        #ggtitle("Observed weekly means and climatology") +
  scale_x_date(breaks = seq(min(all$day.x), max(all$day.x), by = "7 days"),
               minor_breaks = NULL, expand = c(0.03, 0.03),
               label = date_format("%d %b %Y")) +
  theme(axis.text.x = element_text(angle = 45, hjust = 1),
        panel.grid.major = element_line(size = 0.5, linetype = 'solid',
                                        colour = "gray92"),
        panel.grid.minor = element_line(size = 0.25, linetype = 'solid',
                                        colour = "gray92"),
        legend.spacing = unit(-0.2, "cm")) +
  scale_fill_manual(name = NULL,
                    values = c("p10-p90" = cols[3], "p33-p66" = cols[4])) +
  scale_color_manual(name = NULL, values = c("climatological mean" = cols[5],
                                             "observed daily mean" = "grey20",
                                             "observed weekly mean" = "black")) +
  scale_linetype_manual(name = NULL, values = c("climatological mean" = "solid",
                                                "observed daily mean" = "dashed",
                                                "observed weekly mean" = "solid"),
               guide = guide_legend(override.aes = list(lwd = c(0.7, 0.2, 0.7)))) +
    guides(fill = guide_legend(order = 1))
p
ggsave(paste0(output_dir, "CaseStudy_observed_weekly_means_FMI_", year, "_",
       paste0(months, collapse = "-"), ".jpg"), height = 4, width = 6, dpi = 300) 
