library(ncdf4)
library(raster)
library(dplyr)
library(sf)
library(mapview)
library(terra)

# PARAMETERS ----
year <- 2024
home_ada <- "/home/abarrant"
input_stats_obs <- paste0(home_ada,"/esarchive/scratch/ccarnere/RScripts/eval_gene/", year, "/statistics_at_stations.RDS")
input_stats_file <- paste0(home_ada,"/esarchive/scratch/ccarnere/diagnosi/", year, "/statistics_", year, ".nc")
output_shp <- paste0(home_ada, "/esarchive/scratch/abarrant/test_netcdf/statistics_with_obs_", year, ".shp")
output_nc <- paste0(home_ada, "/esarchive/scratch/abarrant/test_netcdf/statistics_with_obs_", year, ".nc")
lambert_crs <- "+proj=lcc +lat_1=43 +lat_2=37 +lat_0=40 +lon_0=-3 +x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs"

# Create output directory
dir.create(dirname(output_shp), showWarnings = FALSE, recursive = TRUE)

# Load data
raster_stack <- rast(input_stats_file)
stats_obs <- readRDS(input_stats_obs)

# Get variable names that match between raster and observations
raster_vars <- names(raster_stack)
obs_vars <- names(stats_obs)
common_vars <- intersect(raster_vars, obs_vars[!obs_vars %in% c("lon", "lat")])

cat("Found", length(common_vars), "matching variables:\n")
cat(paste(common_vars, collapse = ", "), "\n\n")

# Create obs spatial data
obs_sf <- st_as_sf(stats_obs, coords = c("lon", "lat"), crs = 4326)

# Prepare coordinate handling (once for all variables)
obs_coords <- data.frame(
  lon = stats_obs$lon,
  lat = stats_obs$lat
)
valid_coords <- complete.cases(obs_coords)
obs_coords <- obs_coords[valid_coords, ]
obs_points <- SpatialPoints(obs_coords, proj4string = CRS("+proj=longlat +datum=WGS84"))
obs_points_transformed <- spTransform(obs_points, CRS(lambert_crs))

# Initialize lists to store results
modified_rasters <- list()
processed_obs_data <- list()

# Process each common variable
for (var_name in common_vars) {
  cat("Processing variable:", var_name, "\n")
  
  # Extract raster layer
  terra_layer <- raster_stack[[var_name]]
  raster_layer <- raster(terra_layer)
  
  # Create modified copy
  raster_modified <- raster_layer
  
  # Extract cell numbers for observation points
  cell_numbers <- cellFromXY(raster_modified, coordinates(obs_points_transformed))
  
  # Remove cells outside raster domain
  valid_cells <- !is.na(cell_numbers) & valid_coords
  cell_numbers <- cell_numbers[!is.na(cell_numbers)]
  
  # Get replacement values from observations
  replacement_values <- stats_obs[[var_name]][valid_coords][!is.na(cell_numbers)]
  replacement_values <- as.numeric(replacement_values)
  
  # Only replace valid (non-NA/NaN) values
  valid_replacements <- !is.na(replacement_values) & !is.nan(replacement_values)
  cells_to_replace <- cell_numbers[valid_replacements]
  values_to_use <- replacement_values[valid_replacements]
  
  # Replace raster cells with observation values
  raster_modified[cells_to_replace] <- values_to_use
  
  # Store results
  modified_rasters[[var_name]] <- raster_modified
  
  # Prepare observation data for shapefile (clean NAs)
  obs_clean <- obs_sf
  obs_clean[[var_name]] <- as.numeric(obs_clean[[var_name]])
  obs_clean <- obs_clean[!is.na(obs_clean[[var_name]]), ]
  processed_obs_data[[var_name]] <- obs_clean
  
  cat("  - Replaced", length(values_to_use), "cells\n")
}

# Create combined shapefile with all variables
cat("\nCreating combined shapefile...\n")

# Start with the first processed observation dataset
combined_obs <- processed_obs_data[[1]]

# Add other variables by joining (matching by geometry)
if (length(processed_obs_data) > 1) {
  for (i in 2:length(processed_obs_data)) {
    var_name <- names(processed_obs_data)[i]
    temp_data <- processed_obs_data[[i]]
    
    # Create a simple merge based on coordinates (since geometry should be the same)
    # Keep only the variable column and geometry
    temp_data <- temp_data[c(var_name, attr(temp_data, "sf_column"))]
    
    # Merge with combined_obs
    combined_obs <- merge(combined_obs, temp_data, by = attr(combined_obs, "sf_column"), all = TRUE)
  }
}

# Save shapefile
st_write(combined_obs, output_shp, delete_dsn = TRUE)
cat("Shapefile saved to:", output_shp, "\n")

# Create NetCDF with modified rasters
cat("Creating NetCDF file...\n")

# Convert rasters to terra format for NetCDF writing
terra_layers <- list()
for (var_name in names(modified_rasters)) {
  terra_layers[[var_name]] <- rast(modified_rasters[[var_name]])
}

# Combine into a single SpatRaster
combined_raster <- do.call(c, terra_layers)
names(combined_raster) <- names(modified_rasters)

# Write NetCDF
writeCDF(combined_raster, output_nc, overwrite = TRUE)
cat("NetCDF saved to:", output_nc, "\n")

# Optional: Create summary report
cat("\nSummary Report:\n")
cat("===============\n")
cat("Total variables processed:", length(common_vars), "\n")
cat("Variables:", paste(common_vars, collapse = ", "), "\n")
cat("Output shapefile:", output_shp, "\n")
cat("Output NetCDF:", output_nc, "\n")

# Optional: Quick visualization of first variable for verification
if (length(common_vars) > 0) {
  first_var <- common_vars[1]
  cat(paste("\nCreating preview map for", first_var, "...\n"))
  
  # Calculate common color scale
  raster_range <- range(values(modified_rasters[[first_var]]), na.rm = TRUE)
  obs_range <- range(combined_obs[[first_var]], na.rm = TRUE)
  combined_range <- range(c(raster_range, obs_range), na.rm = TRUE)
  
  breaks <- seq(combined_range[1], combined_range[2], length.out = 10)
  colors <- viridisLite::viridis(length(breaks) - 1)
  
  # Create preview map
  m1 <- mapview(modified_rasters[[first_var]], 
                at = breaks,
                col.regions = colors,
                layer.name = paste("Modified", first_var))
  
  m2 <- mapview(combined_obs, 
                zcol = first_var,
                at = breaks,
                col.regions = colors,
                layer.name = paste("Observations", first_var))
  
  preview_map <- m1 + m2
  print(preview_map)
}

cat("\nProcessing completed successfully!\n")