library(ncdf4)
library(lubridate)
library(ggplot2)

# Define location
locations <- list(vriotajo = c(lat = 39.87, lon = 355.87))
location_names <- names(locations)

# Function to find nearest index for latitude/longitude
find_nearest_index <- function(value, grid) {
  which.min(abs(grid - value))
}

# Load temperature data from NetCDF
load_variable_to_cube <- function(var_name, year) {
  all_obs <- list()
  all_days <- c()
  patho <- "/esarchive/recon/ecmwf/era5land/daily/"
  
  start_date <- as.Date(paste0(year, "-03-01"))
  end_date <- as.Date(paste0(year, "-08-31"))
  dates <- seq(start_date, end_date, by = "month")
  
  for (j in seq_along(dates)) {
    date <- format(dates[j], "%Y%m")
    file_path <- paste0(patho, var_name, "/", var_name, "_", date, ".nc")
    cat("Loading:", file_path, "\n")
    
    nc_data <- nc_open(file_path)
    lat_vals <- ncvar_get(nc_data, "lat")
    lon_vals <- ncvar_get(nc_data, "lon")
    time_vals <- ncvar_get(nc_data, "time")
    var_vals <- ncvar_get(nc_data, var_name)
    fill_value <- ncatt_get(nc_data, var_name, "_FillValue")$value
    time_units <- ncatt_get(nc_data, "time", "units")$value
    
    ref_date <- sub("hours since ", "", time_units)
    ref_date <- as.POSIXct(ref_date, tz = "UTC")
    dates_current <- as.Date(ref_date + time_vals * 3600, tz = "UTC")
    dates_current <- format(dates_current, "%Y%m%d")
    
    nc_close(nc_data)
    
    for (i in seq_along(locations)) {
      loc <- locations[[i]]
      lat_idx <- find_nearest_index(loc["lat"], lat_vals)
      lon_idx <- find_nearest_index(loc["lon"], lon_vals)
      
      obs_vals <- var_vals[lon_idx, lat_idx, ]
      obs_vals[obs_vals == fill_value] <- NA
      
      loc_name <- location_names[i]
      all_obs[[loc_name]] <- c(all_obs[[loc_name]], obs_vals)
    }
    
    all_days <- c(all_days, dates_current)
  }
  
  n_loc <- length(locations)
  n_time <- length(all_days)
  result_array <- array(NA, dim = c(n_loc, n_time, 1),
                        dimnames = list(loc = location_names, time = NULL, variable = var_name))
  
  for (i in seq_along(locations)) {
    loc_name <- location_names[i]
    result_array[i, , 1] <- all_obs[[loc_name]]
  }
  
  dates_list <- list(start = as.Date(all_days, format = "%Y%m%d"))
  names(dates_list$start) <- all_days
  
  cube <- list(
    data = result_array,
    Dates = dates_list,
    lat = sapply(locations, function(x) x["lat"]),
    lon = sapply(locations, function(x) x["lon"]),
    var = list(list(name = var_name, units = "K")),
    source = "era5land"
  )
  class(cube) <- "s2dv_cube"
  return(cube)
}

# Compute GDD from first Monday of March to last Sunday of August
compute_gdd <- function(tasmax, tasmin, base_temp = 10) {
  # Compute GDD: average of max and min temp minus base_temp
  gdd_data <- ((tasmax$data + tasmin$data) / 2) - base_temp
  gdd_data[gdd_data < 0] <- 0
  
  all_dates <- as.Date(names(tasmax$Dates$start), format = "%Y%m%d")
  year <- format(all_dates[1], "%Y")
  
  # Define date range from first Monday of March to last Sunday of August
  march_1 <- as.Date(paste0(year, "-03-01"))
  first_monday <- march_1 + ((1 - wday(march_1, week_start = 1)) %% 7)
  
  aug_31 <- as.Date(paste0(year, "-08-31"))
  # Find the last Sunday *on or before* Aug 31
  last_sunday <- aug_31 - ((wday(aug_31) %% 7))  # Same as: wday Sunday = 1
  first_day_of_last_week <- last_sunday - 6  # The Monday of that week
  
  # Filter the time range
  date_mask <- all_dates >= first_monday & all_dates <= last_sunday
  filtered_dates <- all_dates[date_mask]
  filtered_data <- gdd_data[, date_mask, , drop = FALSE]  # Keep 3D
  
  # Ensure full weeks
  n_days <- length(filtered_dates)
  n_weeks <- n_days %/% 7
  remainder <- n_days %% 7
  if (remainder != 0) {
    # Remove trailing days that don’t complete a week
    filtered_data <- filtered_data[, 1:(n_weeks * 7), , drop = FALSE]
    filtered_dates <- filtered_dates[1:(n_weeks * 7)]
  }
  
  # Reshape: [loc, 7, weeks]
  reshaped <- array(filtered_data[, , 1], dim = c(dim(filtered_data)[1], 7, n_weeks))
  weekly_totals <- apply(reshaped, c(1, 3), sum, na.rm = TRUE)
  weekly_cumsum <- t(apply(weekly_totals, 1, cumsum))
  
  # Build result cube
  weekly_dates <- seq(first_monday, by = "1 week", length.out = n_weeks)
  names(weekly_dates) <- format(weekly_dates, "%Y%m%d")
  
  gdd_weekly_cumsum <- list(
    data = array(weekly_cumsum, dim = c(dim(weekly_cumsum), 1)),
    Dates = list(start = weekly_dates),
    lat = tasmax$lat,
    lon = tasmax$lon,
    var = list(list(name = "GDD_cumulative", units = "°C·day")),
    source = tasmax$source
  )
  class(gdd_weekly_cumsum) <- "s2dv_cube"
  return(gdd_weekly_cumsum)
}

# ---- MAIN SCRIPT ----

# Load temperature data for 2024
tasmax_2024 <- load_variable_to_cube("tasmax", 2024)
tasmin_2024 <- load_variable_to_cube("tasmin", 2024)

# Convert from Kelvin to Celsius
tasmax_2024$data <- tasmax_2024$data - 273.15
tasmin_2024$data <- tasmin_2024$data - 273.15

# Compute GDD
gdd_2024 <- compute_gdd(tasmax_2024, tasmin_2024)

# Extract data for plotting
dates_2024 <- as.Date(names(gdd_2024$Dates$start), format = "%Y%m%d")
gdd_values <- gdd_2024$data[1, , 1]

# Create data frame for ggplot
df_2024 <- data.frame(
  Date = dates_2024,
  GDD = gdd_values,
  Year = "2024"
)

# Plot
output_path <- "/esarchive/scratch/ptrascas/R/dev-test_bigpredidata/sunset/sunset/gdd_timeseries.png"

png(filename = output_path, width = 1200, height = 600, res = 150)
ggplot(df_2024, aes(x = Date, y = GDD, color = Year)) +
  geom_line(size = 1.2) +
  labs(title = "Cumulative Weekly Growing Degree Days (GDD)",
       x = "Week",
       y = "Cumulative GDD (°C·day)",
       color = "") +
  scale_x_date(date_labels = "%d-%b", date_breaks = "1 week") +
  theme_minimal(base_size = 14) +
  theme(axis.text.x = element_text(angle = 45, hjust = 1))
dev.off()

