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_file <- paste0(home_ada,"/esarchive/scratch/ccarnere/diagnosi/", year, "/statistics_with_obs_", year, ".nc")
output_file <- paste0(home_ada, "/esarchive/scratch/abarrant/test_netcdf/statistics_with_obs", year, ".tif")
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"
#---


dir.create(dirname(output_file), showWarnings = FALSE, recursive = TRUE)




#get the raster stack with all the layers and select the desired one
raster_stack <- rast(input_stats_file)
terra_conc <- raster_stack[["o3_p93.2"]]  
raster_conc <- raster(terra_conc)             #change it to rasterlayer instead of SpatRaster
#mapview(raster_conc)

#obs stations
stats_obs <- readRDS(input_stats_obs)
obs_sf <- st_as_sf(stats_obs, coords = c("lon", "lat"), crs = 4326)
obs_proj <- st_transform(obs_sf, crs = crs(raster_conc))

#create a copy of the original raster for modification
raster_modified <- raster_conc

#convert observation coordinates to spatial points
obs_coords <- data.frame(
  lon = stats_obs$lon,
  lat = stats_obs$lat
)

#remove rows with nan coordinates (it's not necessary since there are none latlon missing )
valid_coords <- complete.cases(obs_coords)
obs_coords <- obs_coords[valid_coords, ]

#create spatial points (assuming obs coordinates are in WGS84)
obs_points <- SpatialPoints(obs_coords, proj4string = CRS("+proj=longlat +datum=WGS84"))

#transform observation points to the same CRS as the raster
obs_points_transformed <- spTransform(obs_points, CRS(lambert_crs))

#extract cell numbers for the observation points
cell_numbers <- cellFromXY(raster_modified, coordinates(obs_points_transformed))

#remove nan cells (outside the domain of the raster)
valid_cells <- !is.na(cell_numbers) & valid_coords
cell_numbers <- cell_numbers[!is.na(cell_numbers)]


#get the replacement values from stats_obs
replacement_values <- stats_obs$o3_p93.2[valid_coords][!is.na(cell_numbers)]  ### CHANGE DEPENDING ON THE COLUMN
replacement_values <- as.numeric(replacement_values)                          #convert string to numeric values (they come as strings at the stats_obs)


#only replace the cells when the value is not nan or na
valid_replacements <- !is.na(replacement_values) & !is.nan(replacement_values)
cells_to_replace   <- cell_numbers[valid_replacements]              #correct number of cells that dont have nans
values_to_use      <- replacement_values[valid_replacements]        #values to use in the replacement

#replace cells with the values of the obs
raster_modified[cells_to_replace] <- values_to_use

#mapview -----------------------------------------------------------
# raster_range <- range(values(raster_modified), na.rm = TRUE)
# obs_range <- range(as.numeric(obs_sf$o3_p93.2), na.rm = TRUE)
# combined_range <- range(c(raster_range, obs_range), na.rm = TRUE)
# 
# # Create common breaks
# breaks <- seq(combined_range[1], combined_range[2], length.out = 15)
# 
# # Define color palette
# colors <- viridisLite::viridis(length(breaks) - 1)
# #colors <- viridisLite::plasma(length(breaks) - 1)
# 
# m1 <- mapview(raster_modified, at = breaks, col.regions = colors)
# 
# obs_sf$o3_p93.2 <- as.numeric(obs_sf$o3_p93.2)
# obs_sf <- obs_sf[!is.na(obs_sf$o3_p93.2), ]
# m2 <- mapview(obs_sf,  zcol="o3_p93.2", at = breaks, col.regions = colors)
# m1 + m2

#saving new raster with obs replaced
writeRaster(raster_modified, output_file, overwrite = TRUE)

#saving it as a shapefile (fer tmb nc)
polygons <- rasterToPolygons(raster_modified, dissolve=TRUE)      #convert to polygons (spatial points)
polygons_sf <- st_as_sf(polygons)
st_write(polygons_sf, paste0(home_ada, "/esarchive/scratch/abarrant/test_netcdf/statistics_with_obs", year, ".shp"), delete_dsn = TRUE)   #delete_dsn: overwrites any existing file










#optional: summary report------------------------------------------------
cat("\nSummary Report:\n")
cat("===============\n")
cat("Total variables processed:", length(common_vals), "\n")
cat("Variables:", paste(common_vals, 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_vals) > 0) {
  first_var <- common_vals[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")