library(ncdf4)
library(lubridate)
library(ggplot2)
library(startR)
library(reshape2)

# -------------- CONFIG --------------
locations <- list(vriotajo = c(lat = 39.87, lon = 355.87))
location_names <- names(locations)
target_lat <- 39.87
target_lon <- 355.87

output_dir <- "/esarchive/scratch/ptrascas/R/dev-test_bigpredidata/sunset/sunset"
if (!dir.exists(output_dir)) dir.create(output_dir, recursive = TRUE)

year <- 2024
base_temp <- 10
fcst.sdate <- "20240826"
fcst_start_date <- as.Date(fcst.sdate, format = "%Y%m%d")
ltmin <- 1
ltmax <- 35
fcst.nmember <- 16
fcst_path_tasmax <- "/esarchive/exp/ncep/cfs-v2/daily/s2s/tasmax/tasmax_$file_date$.nc"
fcst_path_tasmin <- "/esarchive/exp/ncep/cfs-v2/daily/s2s/tasmin/tasmin_$file_date$.nc"

# ---------- FUNCTIONS ----------
find_nearest_index <- function(value, grid) {
  which.min(abs(grid - value))
}

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-25"))
  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 <- as.POSIXct(sub("hours since ", "", time_units), tz = "UTC")
    dates_current <- as.Date(ref_date + time_vals * 3600)
    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, format(dates_current, "%Y%m%d"))
  }
  
  result_array <- array(NA, dim = c(length(locations), length(all_days), 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]]
  }
  cube <- list(
    data = result_array,
    Dates = list(start = as.Date(all_days, format = "%Y%m%d")),
    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 <- function(tasmax, tasmin, base_temp = 10, initial_gdd = 0, weekly = TRUE) {
  # 1. Daily GDD calculation
  gdd_data <- ((tasmax$data + tasmin$data) / 2) - base_temp
  gdd_data[gdd_data < 0] <- 0
  
  dates <- tasmax$Dates$start
  n_days <- length(dates)
  
  if (weekly) {
    # Trim to full weeks
    n_weeks <- n_days %/% 7
    gdd_data <- gdd_data[, 1:(n_weeks * 7), , drop = FALSE]
    dates <- dates[1:(n_weeks * 7)]
    
    # Reshape to weeks: [location, 7 days, week]
    reshaped <- array(gdd_data[, , 1], dim = c(dim(gdd_data)[1], 7, n_weeks))
    weekly <- apply(reshaped, c(1, 3), sum, na.rm = TRUE)
    
    # Weekly cumulative sum + initial GDD if any
    weekly_cumsum <- t(apply(weekly, 1, cumsum))
    weekly_cumsum <- sweep(weekly_cumsum, 1, initial_gdd, `+`)
    
    result_data <- array(weekly_cumsum, dim = c(dim(weekly_cumsum), 1))
    result_dates <- seq(dates[1], by = "1 week", length.out = n_weeks)
  } else {
    # Daily cumulative sum + initial GDD if any
    daily_cumsum <- t(apply(gdd_data[, , 1, drop = FALSE], 1, cumsum))
    daily_cumsum <- sweep(daily_cumsum, 1, initial_gdd, `+`)
    
    result_data <- array(daily_cumsum, dim = c(dim(daily_cumsum), 1))
    result_dates <- dates
  }
  
  # Build output cube
  cube <- list(
    data = result_data,
    Dates = list(start = result_dates),
    lat = tasmax$lat,
    lon = tasmax$lon,
    var = list(list(name = "GDD_cumulative", units = "°C·day")),
    source = tasmax$source
  )
  class(cube) <- "s2dv_cube"
  return(cube)
}

# ---------- OBSERVED ----------
tasmax_obs <- load_variable_to_cube("tasmax", year)
tasmin_obs <- load_variable_to_cube("tasmin", year)
tasmax_obs$data <- tasmax_obs$data - 273.15
tasmin_obs$data <- tasmin_obs$data - 273.15
gdd_obs <- compute_gdd(tasmax_obs, tasmin_obs, weekly = TRUE)
# Align forecast start date to next Monday after last observation week
last_obs_date <- tail(gdd_obs$Dates$start, 1)
next_monday <- last_obs_date + (8 - wday(last_obs_date)) %% 7
cat("✅ Adjusted forecast start date to next Monday:", as.Date(next_monday), "\n")

# Adjust forecast cube dates accordingly
days_to_skip <- as.numeric(next_monday - fcst_start_date)
if (days_to_skip > 0) {
  tasmax_fc <- tasmax_fc[-seq_len(days_to_skip)]
  tasmin_fc <- tasmin_fc[-seq_len(days_to_skip)]
  fcst_start_date <- next_monday
}
dates_obs <- gdd_obs$Dates$start
gdd_vals_obs <- gdd_obs$data[1, , 1]
last_obs_date <- tail(dates_obs, 1)
final_obs_value <- tail(gdd_vals_obs, 1)
cat("📅 Last weekly observation date:", as.Date(last_obs_date), "\n")
cat("🔢 Last cumulative GDD value:", final_obs_value, "\n")

df_obs <- data.frame(Date = dates_obs, GDD = gdd_vals_obs, Type = "Observed")
df_obs <- data.frame(Date = dates_obs, GDD = gdd_vals_obs, Type = "Observed")

# ---------- FORECAST (1st member only) ----------
file_dates.fcst <- fcst.sdate
dim(file_dates.fcst) <- c(1)

# ---------- FORECAST (1st member only) ----------


fcst_tasmax <- Start(
  dat = fcst_path_tasmax, var = "tasmax",
  file_date = file_dates.fcst,
  time = indices(ltmin:ltmax),
  latitude = values(target_lat),
  longitude = values(target_lon),
  ensemble = indices(1),  # Load only 1st ensemble member
  retrieve = TRUE
)

fcst_tasmin <- Start(
  dat = fcst_path_tasmin, var = "tasmin",
  file_date = file_dates.fcst,
  time = indices(ltmin:ltmax),
  latitude = values(target_lat),
  longitude = values(target_lon),
  ensemble = indices(1),
  retrieve = TRUE
)

# Since only 1 member is loaded, drop unnecessary dims
tasmax_fc <- drop(fcst_tasmax)
tasmin_fc <- drop(fcst_tasmin)
fcst_dates <- seq(fcst_start_date, by = "1 day", length.out = length(tasmax_fc))

tasmax_cube_fc <- list(
  data = array(tasmax_fc - 273.15, dim = c(1, length(tasmax_fc), 1)),
  Dates = list(start = fcst_dates),
  lat = target_lat,
  lon = target_lon,
  var = list(list(name = "tasmax", units = "°C")),
  source = "forecast"
)
class(tasmax_cube_fc) <- "s2dv_cube"

tasmin_cube_fc <- list(
  data = array(tasmin_fc - 273.15, dim = c(1, length(tasmin_fc), 1)),
  Dates = list(start = fcst_dates),
  lat = target_lat,
  lon = target_lon,
  var = list(list(name = "tasmin", units = "°C")),
  source = "forecast"
)
class(tasmin_cube_fc) <- "s2dv_cube"

# Compute forecast GDD with carry-over from observations
gdd_fcst <- compute_gdd(tasmax_cube_fc, tasmin_cube_fc, initial_gdd = final_obs_value, weekly = TRUE)
dates_fcst <- gdd_fcst$Dates$start
gdd_vals_fcst <- gdd_fcst$data[1, , 1]

cat("📅 First forecast weekly date:", as.Date(head(dates_fcst, 1)), "\n")

df_fcst <- data.frame(Date = dates_fcst, GDD = gdd_vals_fcst, Type = "Forecast")


# ---------- PLOT ----------
library(ggplot2)

# Define the output path
output_path <- file.path(output_dir, "GDD_forecast_member1.png")

library(scales)  # for date formatting

g <- ggplot() +
  geom_line(data = df_fcst, aes(x = Date, y = GDD, color = "Forecast"), size = 1) +
  geom_line(data = df_obs, aes(x = Date, y = GDD, color = "Observed"), size = 1.2) +
  labs(title = "Cumulative GDD: Observed vs Forecast (1st Ensemble Member)",
       x = "Week Starting", y = "Cumulative GDD (°C·day)", color = "Type") +
  scale_x_date(date_breaks = "1 week", date_labels = "%d-%b") +
  scale_color_manual(values = c("Observed" = "black", "Forecast" = "blue")) +
  theme_minimal(base_size = 14) +
  theme(axis.text.x = element_text(angle = 45, hjust = 1))  # Rotate for readability

# Save the plot
ggsave(filename = output_path, plot = g, width = 10, height = 5, dpi = 150)

