# Merged and Automated Preprocess Script for NO2 Data and Observational Campaigns
# Author: Andreu Julián Izquierdo
# Date: 02/05/2025

# Load required libraries
library(dplyr)
library(lubridate)
library(sf)
library(readxl)
library(sp)
library(gstat)
library(automap)
library(ggplot2)
library(viridis)

# ---------------------------
# USER PARAMETERS
# ---------------------------
# Base directories for model data and outputs
base_model_in  <- "/esarchive/scratch/ajulian/Aveiro/25m_nmd"
base_model_out <- "/esarchive/scratch/ajulian/Aveiro/25m_nmd_mean"

# Base directory for observational campaign folders
# Each period folder should contain the Excel file named results_<n>th_campaign.xlsx
base_obs_in    <- "/esarchive/scratch/ajulian/Aveiro/Resultats_CAMPANYES"

# Output directory for spatial RDS results
base_out       <- "/esarchive/scratch/ajulian/Aveiro/Scripts"

buildings <- readRDS(file.path(base_out,"buildings_cropped_closer.rds"))
# Periods to process
periods <- 1:6

# Ensure overall output directory exists
dir.create(base_out, recursive = TRUE, showWarnings = FALSE)

# Function to generate ordinal suffix (for filenames)
ordinal_suffix <- function(n) {
  if (n %% 10 == 1 && n %% 100 != 11) return("st")
  if (n %% 10 == 2 && n %% 100 != 12) return("nd")
  if (n %% 10 == 3 && n %% 100 != 13) return("rd")
  return("th")
}

# Coordinate reference systems
crs_utm   <- "+proj=utm +zone=31 +ellps=intl +units=m +no_defs"
crs_wgs84 <- 4326

# ---------------------------
# MAIN LOOP OVER PERIODS
# ---------------------------
# ---------------------------
for (i in periods) {
  period_name <- paste0("Period_", i)
  obs_suffix <- paste0(i, ordinal_suffix(i))
  
  # Create input/output paths for this period
  model_in_dir  <- file.path(base_model_in,  period_name)
  model_out_dir <- file.path(base_model_out, period_name)
  dir.create(model_out_dir, recursive = TRUE, showWarnings = FALSE)
  
  # ----- PART 1: PROCESS MODEL RDS FILES -----
  message("Processing model data for ", period_name)
  rds_files <- list.files(
    path = model_in_dir,
    pattern = "^caliope_urban_results_\\d{8}_NO2\\.rds$",
    full.names = TRUE
  )
  
  if (length(rds_files) == 0) {
    warning("No RDS files found in ", model_in_dir, " – skipping model processing.")
    next
  }
  
  all_daily <- list()
  for (rds_file in rds_files) {
    data_hourly <- readRDS(rds_file)
    data_hourly$date <- as.Date(data_hourly$date)
    
    daily_mean_NO2 <- data_hourly %>%
      group_by(date, X, Y) %>%
      summarise(NO2 = mean(NO2, na.rm = TRUE), .groups = 'drop') %>%
      mutate(date_type = "daily_mean")
    
    data_hourly <- data_hourly %>% mutate(date_type = "hourly")
    
    combined <- bind_rows(data_hourly, daily_mean_NO2)
    out_name <- sub("NO2\\.rds$", "mean_NO2.rds", basename(rds_file))
    saveRDS(combined, file.path(model_out_dir, out_name))
    
    all_daily[[basename(rds_file)]] <- daily_mean_NO2
  }
  
  if (length(all_daily) > 0) {
    merged_daily <- bind_rows(all_daily)
    final_points <- merged_daily %>%
      group_by(X, Y) %>%
      summarise(NO2_mean = mean(NO2, na.rm = TRUE), .groups = 'drop')
    final_sf <- st_as_sf(final_points, coords = c("X", "Y"), crs = crs_utm)
    saveRDS(final_sf, file.path(model_out_dir, paste0("mean_campaign_p", i, ".rds")))
  }
  
  # ----- PART 2: PROCESS OBSERVATIONAL DATA -----
  obs_dirs <- list.dirs(base_obs_in, recursive = FALSE, full.names = TRUE)
  matches  <- obs_dirs[grepl(paste0("^", i), basename(obs_dirs))]
  if (length(matches) == 0) {
    warning("No observational folder found for campaign ", obs_suffix, " – skipping.")
    next
  }
  obs_dir <- matches[1]
  
  message("Processing observational data for campaign ", obs_suffix)
  
  excel_file <- list.files(
    path = obs_dir,
    pattern = paste0("results_", obs_suffix, "_campaign.*\\.xlsx$"),
    full.names = TRUE
  )
  
  if (length(excel_file) != 1) {
    excel_file <- list.files(
      path = obs_dir,
      pattern = "results.*\\.xlsx$",
      full.names = TRUE
    )
  }
  
  if (length(excel_file) != 1) {
    warning("Expected one results Excel in ", obs_dir, "; found ", length(excel_file))
    next
  }
  
  df_obs <- read_excel(excel_file, sheet = 1)
  names(df_obs) <- trimws(names(df_obs))
  df_obs <- df_obs %>%
    rename_with(~"Site", matches("(?i)^site$")) %>%
    rename(pollutant_obs = `Results (ug/m3)`)
  
  sf_select <- df_obs %>%
    select(Site, latitude, longitude) %>%
    st_as_sf(coords = c("longitude", "latitude"), crs = crs_wgs84)
  
  saveRDS(sf_select, file.path(base_out, paste0("sp_coord_dosimeters_p", i, ".rds")))
  saveRDS(df_obs,    file.path(base_out, paste0("df_obs_dosimeters_p", i, ".rds")))
  
  # ----- PART 3: INTERPOLATE MODEL AT OBSERVATION LOCATIONS BY OK -----
  message("Performing OK of model to obs locations for campaign ", obs_suffix)
  
  model_sp <- as(final_sf, "Spatial")
  mod_var_fit <- automap::autofitVariogram(
    formula = NO2_mean ~ 1,
    input_data = model_sp
  )
  
  ok_mod <- gstat(formula = NO2_mean ~ 1, data = model_sp, model = mod_var_fit$var_model)
  
  obs_pts <- as(sf_select, "Spatial")
  obs_pts <- sp::spTransform(obs_pts, CRS(crs_utm))
  
  pred_mod <- predict(ok_mod, newdata = obs_pts)
  
  png(file.path(base_out, paste0("scatter_obs_vs_okmodel_p", i, ".png")), width = 600, height = 600)
  plot(df_obs$pollutant_obs, pred_mod$var1.pred,
       xlab = expression("Observed (" * mu * "g/m"^3 * ")"),  # Add units to the x-axis label
       ylab = expression("Raw Model (" * mu * "g/m"^3 * ")"),  # Add units to the y-axis label
       xlim = c(0, 50), ylim = c(0, 50),
       main = paste("Obs vs OK-Model for campaign", obs_suffix))
  abline(0, 1, col = "red", lty = 2)
  dev.off()
  
  # Convert OK predictions to sf
  ok_sf <- st_as_sf(pred_mod, coords = c("X", "Y"), crs = crs_utm)
  
  # Ensure both layers are in same CRS as buildings
  ok_sf <- st_transform(ok_sf, crs_utm)
  buildings <- st_transform(buildings, crs_utm)  # Redundant if already correct
  
  
  # Bounding box around OK-predicted points
  ok_bbox <- st_bbox(ok_sf)
  margin <- 100
  xlim <- c(ok_bbox["xmin"] - margin, ok_bbox["xmax"] + margin)
  ylim <- c(ok_bbox["ymin"] - margin, ok_bbox["ymax"] + margin)
  
  # Create plot
  png(file.path(base_out, paste0("map_okmodel_p", i, ".png")), width = 600, height = 600)
  print(
    ggplot() +
      geom_sf(data = buildings, fill = "grey90", color = "black", size = 0.2) +
      geom_sf(data = ok_sf, aes(color = var1.pred), size = 2) +
      scale_color_viridis_c(name = expression(OK ~ Model ~ NO[2] ~ "(" * mu * "g/m"^3 * ")")) +
      labs(title = paste("OK-model NO2 predictions – Campaign", i)) +
      coord_sf(xlim = xlim, ylim = ylim) +
      theme_minimal(base_size = 11) +
      theme(
        legend.position = "right",
        legend.title = element_text(size = 9),
        legend.text = element_text(size = 8),
        axis.title = element_blank(),
        axis.text = element_blank(),
        axis.ticks = element_blank(),
        plot.title = element_text(size = 12, face = "bold", hjust = 0.5)
      )
  )
  dev.off()
  
}

