################## Excess heat obs + forecast
# This script aims to obtain a forecast of the
# excess heat accumulated by weeks
# first it shows the observations until a certain date
# then it shows the forecast plume and ensemble mean
# Created by Paloma Trascasa-Castro


Sys.setenv(LANG = "en")
Sys.setlocale("LC_TIME","en_GB")
library(startR)
library(s2dverification)
library(abind)
library(CSTools)
library(zeallot)
library(s2dv)
library(easyVerification)
library(ClimProjDiags)
library(ggplot2)
library(CSIndicators)

# Test for 2022 just because there's data 
# Now I'm just using mean daily temperature but the average of tmin and tmax should be taken instead

path <- "/esarchive/recon/ecmwf/era5land/daily_mean/tas_f1h/tas_202205.nc"
path <- paste0(path)

# Select 1 point
Granollers <- c(lat = 41.61, lon = 2.3)

day <- Start(dat = path,
                  var= 'tas',
                  lat = values(Granollers[[1]]),
                  lon = values(Granollers[[2]]),
                  time = 'all',
                  retrieve = T)
                  as.vector(attributes(day)$Variables$common$lat)
                  as.vector(attributes(day)$Variables$common$lon)

# Now I have an array with just 1 dimension: time

# Convert to Celsius
day <- as.s2dv_cube(day)

day$data = day$data - 273.15

# If threshold is 10C, calculate exceeding temperature values and accumulate.

# Step 1, calculate excess heat above 10C and remove anything below 10

day$data <- day$data - 10

# Replace negative anomalies with 0
day$data[day$data < 0] <- 0

# Now accumulate degrees:
# object needs to be an s2dv_cube

day$data <- cumsum(day$data)

dev.new()
matplot(1:31,day$data, type = 'l', lwd = 2)

############## FORECAST ##################

path_for <- "/esarchive/exp/ecmwf/system5c3s/daily_mean/tas_f6h/tas_20220601.nc"
path_for <- paste0(path_for)


forecast <- Start(dat = path_for,
                  var= 'tas',
                  lat = values(Granollers[[1]]),
                  lon = values(Granollers[[2]]),
                  time = 1:7,
		  ensemble = 'all',
                  retrieve = T)

                  as.vector(attributes(day)$Variables$common$lat)
                  as.vector(attributes(day)$Variables$common$lon)


# Convert to Celsius
forecast <- as.s2dv_cube(forecast)

forecast$data = forecast$data - 273.15

# If threshold is 10C, calculate exceeding temperature values and accumulate.

# Step 1, calculate excess heat above 10C and remove anything below 10

forecast$data <- forecast$data - 10

# Replace negative anomalies with 0
forecast$data[forecast$data < 0] <- 0

# Now accumulate degrees along the "time" dimension:

# Take the last value of observations, add it as the first timestep of the forecast and accumulate from there
lastday <- day$data[31]

# Example dimensions
dat <- 1
var <- 1
lat <- 1
lon <- 1
time <- 7
ensemble <- 51

# Create example data (assuming forecast$data is a numeric array)

# Define the new timestep value
new_timestep_value <- lastday

# Create a new array with extended dimensions
new_dim <- c(dat, var, lat, lon, time + 1, ensemble)
new_data <- array(NA, dim = new_dim)

# Set the new timestep value at the beginning
new_data[, , , , 1, ] <- new_timestep_value

# Copy existing data to the new array
new_data[, , , , 2:(time + 1), ] <- forecast$data

# Replace the original data with the new array
forecast$data <- new_data

forecast$data <- apply(forecast$data, c(1,2,3,4,6),cumsum)

values <- matrix(forecast$data[, 1, 1, 1, 1, ], nrow = 8, ncol = 51)

# Set up the plot
plot(1:31, day$data, type = "l", lwd = 2, col = "blue",
     xlab = 'Timesteps', ylab = 'Values', xlim = c(1, 39), ylim = range(c(day$data, values)),
     main = 'Combined Plot of day$data and values')

# Add the 'values' data as a contiguous time series from timestep 32 onwards
matlines(31:38, values, type = 'l', col = 1:ncol(values))

# Adding legend for the plots
legend('topright', legend = c("day$data", paste('member', 1:ncol(values))),
       col = c("blue", 1:ncol(values)), lty = 1, cex = 0.8)