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

# PARAMETERS ----
year <- 2024
local_script <- "/home/abarrant" 

input_stats_obs <- paste0(local_script,"/esarchive/scratch/ccarnere/RScripts/eval_gene/", year, "/statistics_at_stations.RDS")
input_stats_file <- paste0(local_script,"/esarchive/scratch/ccarnere/diagnosi/", year, "/statistics_", year, ".nc")
output_shp <- paste0(local_script, "/esarchive/scratch/abarrant/test_netcdf/statistics_with_obs_", year, ".shp")
output_nc <- paste0(local_script, "/esarchive/scratch/abarrant/test_netcdf/statistics_with_obs_", year, ".nc")
output_gpkg <- paste0(local_script, "/esarchive/scratch/abarrant/test_netcdf/statistics_with_obs_", year, ".gpkg")

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

# EXTRACT CRS FROM ORIGINAL FILE -----------------------------------------------

nc_orig <- nc_open(input_stats_file)
crs_info <- NULL
lambert_crs <- NULL

#copy existing CRS variable
if("crs" %in% names(nc_orig$var) || "Lambert_conformal" %in% names(nc_orig$var)) {
  crs_var_name <- ifelse("crs" %in% names(nc_orig$var), "crs", "Lambert_conformal")
  
  #extract CRS attributes from original file
  crs_atts <- ncatt_get(nc_orig, crs_var_name)
  
  if("grid_mapping_name" %in% names(crs_atts) && crs_atts$grid_mapping_name == "lambert_conformal_conic") {
    #extract Lambert Conformal Conic parameters
    std_par1 <- ifelse("standard_parallel" %in% names(crs_atts), as.numeric(strsplit(crs_atts$standard_parallel, ",")[[1]][1]), 43)
    std_par2 <- ifelse("standard_parallel" %in% names(crs_atts), as.numeric(strsplit(crs_atts$standard_parallel, ",")[[1]][2]), 37)
    lat_0 <- ifelse("latitude_of_projection_origin" %in% names(crs_atts), crs_atts$latitude_of_projection_origin, 40)
    lon_0 <- ifelse("longitude_of_central_meridian" %in% names(crs_atts), crs_atts$longitude_of_central_meridian, -3)
    false_easting <- ifelse("false_easting" %in% names(crs_atts), crs_atts$false_easting, 0)
    false_northing <- ifelse("false_northing" %in% names(crs_atts), crs_atts$false_northing, 0)
    
    lambert_crs <- paste0("+proj=lcc +lat_1=", std_par1, " +lat_2=", std_par2, 
                          " +lat_0=", lat_0, " +lon_0=", lon_0, 
                          " +x_0=", false_easting, " +y_0=", false_northing, 
                          " +datum=WGS84 +units=m +no_defs")
    crs_info <- crs_atts
  }
}

#Check global attributes
if(is.null(crs_info)) {
  global_atts <- ncatt_get(nc_orig, 0)  # 0 for global attributes
  if("crs" %in% names(global_atts) || "projection" %in% names(global_atts)) {
    cat("Found CRS info in global attributes\n")
  }
}

#original definition if nothing found
if(is.null(lambert_crs)) {
  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"
  cat("Using fallback CRS definition\n")
}

cat("Using CRS:", lambert_crs, "\n")

# REPLACING OBS IN RASTER ------------------------------------------------------

#load data with terra library
raster_stack <- rast(input_stats_file)
raster_vals <- names(raster_stack)

#obs
stats_obs <- readRDS(input_stats_obs)
obs_vals <- names(stats_obs)
common_vals <- intersect(raster_vals, obs_vals[!obs_vals %in% c("lon", "lat")])
cat(paste0("Common values: ", paste(common_vals, collapse = ", ")), "\n\n")

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

#prepare coordinate handling
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))



#create netcdf with proper CRS -------------------------------------------------

var_names <- setdiff(names(raster_stack), c("lat", "lon", "Lambert_conformal", "crs"))
template <- rast(input_stats_file, subds = var_names[1])

#set the CRS properly
crs(template) <- lambert_crs

#get coordinate values from the template (these should match original file)
xvals <- xFromCol(template, 1:ncol(template))
yvals <- yFromRow(template, nrow(template):1)  #flip for NetCDF convention! or not if its already right

x_dim_name <- "x"
y_dim_name <- "y"

x_dim <- ncdim_def(x_dim_name, "m", xvals)
y_dim <- ncdim_def(y_dim_name, "m", yvals)
crs_var <- ncvar_def("crs", "", NULL, prec = "char")


var_defs <- list()
for (var_name in var_names) {
  #extract units from original NetCDF
  units <- "unknown"
  if (var_name %in% names(nc_orig$var)) {
    if("units" %in% names(nc_orig$var[[var_name]]$atts)) {
      units <- nc_orig$var[[var_name]]$atts$units
    } else {
      units_att <- ncatt_get(nc_orig, var_name, "units")
      if(units_att$hasatt) units <- units_att$value
    }
  }
  
  var_def <- ncvar_def(var_name, units, list(x_dim, y_dim), missval = NA, prec = "float")
  var_defs[[var_name]] <- var_def
}

#create NetCDF with all variables + CRS
all_vars <- c(var_defs, list(crs = crs_var))
nc_out <- nc_create(output_nc, all_vars)

#CRS metadata
ncatt_put(nc_out, "crs", "grid_mapping_name", "lambert_conformal_conic")
ncatt_put(nc_out, "crs", "standard_parallel", c(37.0, 43.0))
ncatt_put(nc_out, "crs", "latitude_of_projection_origin", 40.0)
ncatt_put(nc_out, "crs", "longitude_of_central_meridian", -3.0)
ncatt_put(nc_out, "crs", "false_easting", 0.0)
ncatt_put(nc_out, "crs", "false_northing", 0.0)
ncatt_put(nc_out, "crs", "semi_major_axis", 6378137.0)
ncatt_put(nc_out, "crs", "inverse_flattening", 298.257223563)
ncatt_put(nc_out, "crs", "proj4_params", lambert_crs)

#add global attributes
ncatt_put(nc_out, 0, "Conventions", "CF-1.8")
ncatt_put(nc_out, 0, "title", paste("Statistics with observations", year))
ncatt_put(nc_out, 0, "institution", "Your Institution")
ncatt_put(nc_out, 0, "source", "Merged raster and observation data")

#copy other global attributes from original file if they exist
orig_global_atts <- ncatt_get(nc_orig, 0)
for(att_name in names(orig_global_atts)) {
  if(!att_name %in% c("Conventions", "title", "institution", "source")) {
    ncatt_put(nc_out, 0, att_name, orig_global_atts[[att_name]])
  }
}

#close originL nc file
nc_close(nc_orig)

#function to properly orient matrix for NetCDF
orient_for_netcdf <- function(raster_layer) {
  mat <- as.matrix(raster_layer, wide = TRUE)
  
  #flip vertically to match NetCDF convention (bottom-to-top)
  t(mat[nrow(mat):1, ]) #ifnot remove [nrow(mat):1, ] to not flip it!
}

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

#process each common variable
for (var_name in common_vals) {
  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))
  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) 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
  cat("  - Replaced", length(values_to_use), "cells\n")
  
  # Write data to NetCDF with proper orientation
  ncvar_put(nc_out, var_name, orient_for_netcdf(raster_modified))
  ncatt_put(nc_out, var_name, "grid_mapping", "crs")
}

nc_close(nc_out)


# CREATE GEOPACKAGE VERSION (open as vector in QGIS) ---------------------------
cat("\nCreating GeoPackage version...\n")

spat_rasters <- lapply(modified_rasters, function(x) rast(x))
spat_stack <- rast(spat_rasters)


# # Convert to sf for geopackage (this will create points) -> uncoment if you want points
# # Option 1: Convert raster to points
# raster_points <- as.data.frame(spat_stack, xy = TRUE)
# raster_points <- raster_points[complete.cases(raster_points), ]
# 
# if(nrow(raster_points) > 0) {
#   points_sf <- st_as_sf(raster_points, coords = c("x", "y"), crs = lambert_crs)
#   st_write(points_sf, output_gpkg, layer = "raster_points", delete_dsn = TRUE)
#   cat("Saved raster as points in GeoPackage\n")
# }

# Option 2: Convert raster to polygons (warning: this can be memory intensive) 
for(i in 1:nlyr(spat_stack)) {
  layer_name <- names(spat_stack)[i]
  if(is.null(layer_name) || layer_name == "") {
    layer_name <- paste0("layer_", i)
  }
  
  raster_polygons <- as.polygons(spat_stack[[i]])
  polygons_sf <- st_as_sf(raster_polygons)
  
  if(i == 1) {
    st_write(polygons_sf, output_gpkg, layer = layer_name, delete_dsn = TRUE)
  } else {
    st_write(polygons_sf, output_gpkg, layer = layer_name, append = TRUE)
  }
  
  cat("Saved layer", i, "as polygons:", layer_name, "\n")
}



cat("Process completed successfully!\n")
cat("NetCDF file:", output_nc, "\n")
cat("GeoPackage file:", output_gpkg, "\n")


#check reopening nc saved file ------------------------------------------------
raster_check <- rast(output_nc)

t_layer <- raster_check[[6]]
r_layer <- raster(t_layer)
mapview(r_layer)





