####################################################################################
#                                                                                  #
#    Script for Downscaling Air Quality Model Data                                 #
#                                                                                  #
#    This script performs the following steps:                                     #
#    1. Load observation and model data                                            #
#    2. Interpolate raw model data to observation locations                        #
#    3. Prepare data for Universal Kriging                                         #
#    4. Universal Kriging and LOOCV                                                #       
#    5. Collect and Save Summary Statistics                                        #
#    6. Analysis of the results                                                    #
#                                                                                  #
#################################################################################### 

# Load necessary libraries
library(sf)
library(sp)
library(gstat)
library(automap)
library(mapview)
library(ggplot2)
library(openair)
library(dplyr)
library(tidyr)
library(grid)
library(gridExtra)

# Set directories
wdir <- "/esarchive/scratch/ajulian/scripts/Open_Data/Downscaling_kriging/Data/"
wdirpredictor <-  "/esarchive/scratch/ajulian/scripts/Open_Data/Traffic_MPCS/Results/"
wdir2 <- "/esarchive/scratch/ajulian/scripts/Open_Data/Traffic_buffers/Data/"
wwind <- "/gpfs/scratch/bsc32/bsc032604/Rscripts/Downscaling/Results"
wout <- "/esarchive/scratch/ajulian/scripts/Open_Data/Downscaling_kriging/Results/"

# Set projections
crs_latlon <- "+proj=longlat +datum=WGS84 +ellps=WGS84 +towgs84=0,0,0" 
crs_utm    <- "+proj=utm +zone=31 +ellps=intl +units=m +no_defs"      

####################################################################################
# Step 1: Load observation and model data                                          #
####################################################################################

sp_model <- readRDS(paste0(wdir, "annual_NO2_at_AQMS.RDS"))                                          # Model at AQMS
sp_obs_locations <- readRDS(paste0(wdir2,"stations_sf.RDS"))                                         # Obs stations coordinates
df_obs <- readRDS(paste0(wdir,"statistics_at_stations.RDS"))                                         # Obs data
traffic_buffers <- readRDS(file.path(wdirpredictor, "traffic_normalized_buffers_AQMS_combined.RDS")) # Traffic buffers
wind_traffic_buffers <- readRDS(file.path(wdirpredictor, "Wind_traffic_buffers_AQMS_combined.RDS"))  # Wind Traffic buffers

# Convert sp_obs_locations to sf and transform CRS
sp_obs_locations_sf <- st_as_sf(sp_obs_locations, coords = c("lon", "lat"), crs = crs_latlon)
sp_obs_locations_sf <- st_transform(sp_obs_locations_sf, st_crs(crs_utm))
traffic_buffers <- st_transform(traffic_buffers, st_crs(crs_utm))
wind_traffic_buffers <- st_transform(wind_traffic_buffers, st_crs(crs_utm))
sf_model <- st_as_sf(sp_model, coords = c("lon", "lat"), crs = crs_latlon)
sf_model <- st_transform(sp_model, st_crs(crs_utm))
sp_model <- as(sf_model, "Spatial")

df_obs <- df_obs %>%
  filter(!is.na(.[[4]]))

####################################################################################
# Step 2: Interpolate raw model data to observation locations                      #
####################################################################################

# Variogram fitting for pollutant
pollutant.fit_ok <- automap::autofitVariogram(formula = NO2_mean ~ 1, input_data = sp_model)

plot(pollutant.fit_ok)

# Create gstat object for kriging
ok_gstat <- gstat(formula = NO2_mean ~ 1, data = sp_model, model = pollutant.fit_ok$var_model, nmax = 25)

# Interpolation using Ordinary Kriging at observation locations
OK_obs <- predict(ok_gstat, newdata = sp_obs_locations_sf)

####################################################################################
# Step 3: Prepare data for Universal Kriging                                       #
####################################################################################

df_model_traffic <- sp_obs_locations_sf %>%
  st_drop_geometry() %>%
  transmute(
    lat = st_coordinates(sp_obs_locations_sf)[, 2],
    lon = st_coordinates(sp_obs_locations_sf)[, 1],
    station_type = station_type,
    station_area = station_area,
    code = code,
    pollutant_model_pred = OK_obs$var1.pred
  ) %>%
  # Join with traffic buffers
  inner_join(
    traffic_buffers %>% 
      st_drop_geometry() %>% 
      select(code, starts_with("buffer_"), traffic_100_300_500),
    by = "code"
  ) %>%
  # Join with wind traffic buffers
  inner_join(
    wind_traffic_buffers %>% 
      st_drop_geometry() %>% 
      select(code, starts_with("buffer_"), traffic_100_300_500),
    by = "code",
    suffix = c("_traffic", "_wind")
  )

df_model_traffic <- df_model_traffic %>%
  transmute(
    lat, lon, ref_obs = code, pollutant_model_pred, station_type, station_area,
    buffer_value_5m = buffer_5m_normalized_traffic,
    buffer_value_10m = buffer_10m_normalized_traffic,
    buffer_value_15m = buffer_15m_normalized_traffic,
    buffer_value_25m = buffer_25m_normalized_traffic,
    buffer_value_50m = buffer_50m_normalized_traffic,
    buffer_value_100m = buffer_100m_normalized_traffic,
    buffer_value_300m = buffer_300m_normalized_traffic,
    buffer_value_500m = buffer_500m_normalized_traffic,
    buffer_value_1000m = buffer_1000m_normalized_traffic,
    buffer_value_2000m = buffer_2000m_normalized_traffic,
    buffer_value_5000m = buffer_5000m_normalized_traffic,
    buffer_value_10000m = buffer_10000m_normalized_traffic,
    buffer_value_combined = traffic_100_300_500_traffic,
    wind_buffer_value_5m = buffer_5m_normalized_wind,
    wind_buffer_value_10m = buffer_10m_normalized_wind,
    wind_buffer_value_15m = buffer_15m_normalized_wind,
    wind_buffer_value_25m = buffer_25m_normalized_wind,
    wind_buffer_value_50m = buffer_50m_normalized_wind,
    wind_buffer_value_100m = buffer_100m_normalized_wind,
    wind_buffer_value_300m = buffer_300m_normalized_wind,
    wind_buffer_value_500m = buffer_500m_normalized_wind,
    wind_buffer_value_1000m = buffer_1000m_normalized_wind,
    wind_buffer_value_2000m = buffer_2000m_normalized_wind,
    wind_buffer_value_5000m = buffer_5000m_normalized_wind,
    wind_buffer_value_10000m = buffer_10000m_normalized_wind,
    wind_buffer_value_combined = traffic_100_300_500_wind
  )

# Convert to sf object
sf_model_traffic <- st_as_sf(df_model_traffic, coords = c("lon", "lat"), crs = st_crs(sp_obs_locations_sf))
sp_model_traffic <- as(sf_model_traffic, "Spatial")

sf_model_traffic <- sf_model_traffic %>%
  mutate(
    lon = st_coordinates(sf_model_traffic)[, 1],
    lat = st_coordinates(sf_model_traffic)[, 2]
  )

# Create sf_model_obs by merging df_obs with sf_model_traffic by the code
sf_model_obs <- df_obs %>%
  rename(ref_obs = station.code) %>%  
  select(ref_obs, pollutant_obs = no2_annual_mean) %>%  
  mutate(pollutant_obs = as.numeric(pollutant_obs)) %>%  
  inner_join(sf_model_traffic %>% st_drop_geometry(), by = "ref_obs") %>%
  st_as_sf(coords = c("lon", "lat"), crs = st_crs(sp_obs_locations_sf))

sp_model_obs <- as(sf_model_obs, "Spatial")

####################################################################################
# Step 4: Universal Kriging and LOOCV                                              #
####################################################################################

# Function to perform Universal Kriging
perform_universal_kriging <- function(formula, input_data, combined_data, verbose = TRUE, pdf_path = file.path(wout, "Variograms.pdf")) {
  # Fit the variogram
  variogram_fit <- automap::autofitVariogram(formula = formula, input_data = input_data, verbose = verbose)
  
  # Create gstat model
  gstat_model <- gstat(formula = formula, data = input_data, model = variogram_fit$var_model)
  
  # Predict using kriging
  refined_corrected <- predict(gstat_model, newdata = combined_data)

  return(list(
    prediction = refined_corrected,
    variogram_fit = variogram_fit,
    formula = formula
  ))
}

# Function to perform LOOCV
perform_loocv <- function(formula, input_data, variogram_model, nmax = 25) {
  loocv_result <- krige.cv(formula = formula, input_data, model = variogram_model)
  return(loocv_result)
}

# Define Universal Kriging formulas
formulas_uk <- list(
  pollutant_obs ~ pollutant_model_pred,
  pollutant_obs ~ pollutant_model_pred + buffer_value_5m,
  pollutant_obs ~ pollutant_model_pred + buffer_value_10m,
  pollutant_obs ~ pollutant_model_pred + buffer_value_15m,
  pollutant_obs ~ pollutant_model_pred + buffer_value_25m,
  pollutant_obs ~ pollutant_model_pred + buffer_value_50m,
  pollutant_obs ~ pollutant_model_pred + buffer_value_100m,
  pollutant_obs ~ pollutant_model_pred + buffer_value_300m,
  pollutant_obs ~ pollutant_model_pred + buffer_value_500m,
  pollutant_obs ~ pollutant_model_pred + buffer_value_1000m,
  pollutant_obs ~ pollutant_model_pred + buffer_value_2000m,
  pollutant_obs ~ pollutant_model_pred + buffer_value_5000m,
  pollutant_obs ~ pollutant_model_pred + buffer_value_10000m,
  pollutant_obs ~ pollutant_model_pred + wind_buffer_value_5m,
  pollutant_obs ~ pollutant_model_pred + wind_buffer_value_10m,
  pollutant_obs ~ pollutant_model_pred + wind_buffer_value_15m,
  pollutant_obs ~ pollutant_model_pred + wind_buffer_value_25m,
  pollutant_obs ~ pollutant_model_pred + wind_buffer_value_50m,
  pollutant_obs ~ pollutant_model_pred + wind_buffer_value_100m,
  pollutant_obs ~ pollutant_model_pred + wind_buffer_value_300m,
  pollutant_obs ~ pollutant_model_pred + wind_buffer_value_500m,
  pollutant_obs ~ pollutant_model_pred + wind_buffer_value_1000m,
  pollutant_obs ~ pollutant_model_pred + wind_buffer_value_2000m,
  pollutant_obs ~ pollutant_model_pred + wind_buffer_value_5000m,
  pollutant_obs ~ pollutant_model_pred + wind_buffer_value_10000m,
  pollutant_obs ~ pollutant_model_pred + buffer_value_100m + buffer_value_300m + buffer_value_500m,
  pollutant_obs ~ pollutant_model_pred + wind_buffer_value_100m + wind_buffer_value_300m + wind_buffer_value_500m
)

# Define model names
model_names <- c(
  "uk_base",
  "uk_buffer_5m", "uk_buffer_10m", "uk_buffer_15m", "uk_buffer_25m", "uk_buffer_50m", "uk_buffer_100m","uk_buffer_300m",
  "uk_buffer_500m", "uk_buffer_1000m", "uk_buffer_2000m", "uk_buffer_5000m", "uk_buffer_10000m",
  "uk_wind_buffer_5m", "uk_wind_buffer_10m", "uk_wind_buffer_15m", "uk_wind_buffer_25m", "uk_wind_buffer_50m", "uk_wind_buffer_100m",
  "uk_wind_buffer_300m", "uk_wind_buffer_500m", "uk_wind_buffer_1000m", "uk_wind_buffer_2000m", "uk_wind_buffer_5000m", "uk_wind_buffer_10000m",
  "uk_combined_circular_100m_300m_500m", "uk_combined_wind_100m_300m_500m"
)

# Perform Universal Kriging for all formulas
uk_results <- lapply(formulas_uk, function(f) {
  perform_universal_kriging(formula = f, input_data = sp_model_obs, combined_data = sp_model_traffic)
})

# Perform LOOCV for all models
loocv_results <- lapply(uk_results, function(res) {
  perform_loocv(
    formula = res$formula,
    input_data = sp_model_obs,
    variogram_model = res$variogram_fit$var_model
  )
})

#Visualization of uk results

model_26_prediction <- uk_results[[26]]$prediction

mapview(model_26_prediction, zcol = "var1.pred", layer.name = "Model 26 Predictions")

####################################################################################
# Step 5: Collect and Save Summary Statistics                                      #
####################################################################################


stats_results <- data.frame()

for (i in 1:length(loocv_results)) {
  loocv_data <- data.frame(
    obs = loocv_results[[i]]$observed,
    pred = loocv_results[[i]]$var1.pred,
    raw = sf_model_obs$pollutant_model_pred
  )
  
  stats <- modStats(loocv_data, mod = "pred", obs = "obs")
  stats$model <- model_names[i]
  stats_results <- rbind(stats_results, stats)
}

# Save the summary statistics
write.csv(stats_results, file = paste0(wout, "stats_results_with_combined_buffers_2024.csv"), row.names = FALSE)
pdf(file = paste0(wout, "stats_results_2024_traffic_wind.pdf"), width = 15, height = 15)
grid::grid.newpage()
gridExtra::grid.table(stats_results)
dev.off()

####################################################################################
# Step 6: Analysis of the results                                                  #
####################################################################################


# Add ref_obs to loocv_data by matching coordinates
loocv_data <- loocv_data %>%
  mutate(ref_obs = sf_model_traffic$ref_obs[match(
    paste(obs, pred),  # Use obs and pred as placeholders for coordinates
    paste(sf_model_traffic$pollutant_model_pred, sf_model_traffic$pollutant_model_pred)  # Adjust as needed
  )])


# Merge station_type into loocv_data
loocv_data <- merge(
  loocv_data,
  sf_model_traffic %>% st_drop_geometry() %>% select(ref_obs, station_type),
  by =  "ref_obs"
)

#########
# Table bias per type and area of AQMS
#########

compute_bias_summary_group <- function(loocv_results, sf_model_traffic, model_names, grouping_var) {
  bias_summary_long <- data.frame()

  for (i in seq_along(loocv_results)) {
    loocv_df <- as.data.frame(loocv_results[[i]])

    # Convert to sf and spatially join
    loocv_sf <- st_as_sf(loocv_df, coords = c("coords.x1", "coords.x2"), crs = st_crs(sf_model_traffic))
    loocv_sf <- st_join(loocv_sf, sf_model_traffic[, c("ref_obs", grouping_var)], join = st_nearest_feature)
    loocv_df <- st_drop_geometry(loocv_sf)

    # Skip NA predictions or observations
    loocv_df <- loocv_df %>% filter(!is.na(var1.pred) & !is.na(observed))

    # Standardize station_type if needed
    if (grouping_var == "station_type") {
      loocv_df <- loocv_df %>%
        mutate(station_type = case_when(
          station_type %in% c("Trànsit", "T") ~ "T",
          station_type %in% c("Fons", "F") ~ "F",
          station_type %in% c("Industrial", "I") ~ "I",
          TRUE ~ station_type
        ))
    }

    # Bias direction
    loocv_df <- loocv_df %>%
      mutate(bias = ifelse(var1.pred > observed, "Over", "Under"))

    # Group and summarize
    stats <- loocv_df %>%
      group_by(.data[[grouping_var]], bias) %>%
      summarise(n = n(), .groups = "drop") %>%
      group_by(.data[[grouping_var]]) %>%
      mutate(pct = round(100 * n / sum(n), 1)) %>%
      ungroup() %>%
      mutate(Model = model_names[i]) %>%
      rename(group_value = .data[[grouping_var]])

    bias_summary_long <- bind_rows(bias_summary_long, stats)
  }

  return(bias_summary_long)
}

# Compute both summaries
bias_type <- compute_bias_summary_group(loocv_results, sf_model_traffic, model_names, "station_type")
bias_area <- compute_bias_summary_group(loocv_results, sf_model_traffic, model_names, "station_area")

# Process each summary into wide format
process_wide_summary <- function(df, group_label) {
  df <- df %>%
    mutate(column = paste0(group_value, " [", group_label, "]: % ", bias)) %>%
    select(Model, column, pct) %>%
    pivot_wider(names_from = column, values_from = pct, values_fill = 0)
  return(df)
}

wide_type <- process_wide_summary(bias_type, "Type")
wide_area <- process_wide_summary(bias_area, "Area")

# Merge the two summaries
bias_summary_merged <- full_join(wide_type, wide_area, by = "Model") %>%
  select(Model, sort(names(.)[-1]))  # keep Model first, sort others

# Remove "[Type]" from column names
colnames(wide_type) <- gsub("\\s*\\[Type\\]", "", colnames(wide_type))

# Remove "[Area]" from column names
colnames(wide_area) <- gsub("\\s*\\[Area\\]", "", colnames(wide_area))


# Export each summary to its own page in the same PDF
pdf(file = paste0(wout, "bias_summary_type_area_2024.pdf"), width = 15, height = 10)

# Page 1: station_type
grid.newpage()
grid.text("Bias Summary by Station Type", gp = gpar(fontsize = 20, fontface = "bold"))
grid.newpage()
grid.table(as.data.frame(wide_type))

# Page 2: station_area
grid.newpage()
grid.text("Bias Summary by Station Area", gp = gpar(fontsize = 20, fontface = "bold"))
grid.newpage()
grid.table(as.data.frame(wide_area))

dev.off()
# Combined CSV with all columns
write.csv(bias_summary_merged, file = paste0(wout, "bias_summary_type_area_2024.csv"), row.names = FALSE)


##########
# Loocv scatter plot
##########

cv_plot <- function(loocv_data, model_name = "Universal Kriging") {
  # Define colors for station types
  unique_station_types <- unique(loocv_data$station_type)
  station_type_colors <- setNames(
    rainbow(length(unique_station_types)),  # Generate distinct colors
    unique_station_types
  )
  
  # Define shapes for station areas
  unique_station_areas <- unique(loocv_data$station_area)
  station_area_shapes <- setNames(
    seq(15, 15 + length(unique_station_areas) - 1),  # Generate distinct shapes
    unique_station_areas
  )
  
  # Plot observed vs predicted with colors for station_type and shapes for station_area
  plot(loocv_data$observed, loocv_data$var1.pred, 
       xlab = expression("Observed (" * mu * "g/m"^3 * ")"),
       ylab = expression("Predicted (" * mu * "g/m"^3 * ")"),
       main = model_name,
       xlim = c(0, max(loocv_data$observed, loocv_data$var1.pred, na.rm = TRUE)), 
       ylim = c(0, max(loocv_data$observed, loocv_data$var1.pred, na.rm = TRUE)),
       col = station_type_colors[loocv_data$station_type],  # Color points by station_type
       pch = station_area_shapes[loocv_data$station_area],  # Shape points by station_area
       cex = 1.5)  # Increase point size
  
  # Add error bars
  arrows(x0 = loocv_data$observed, 
         y0 = loocv_data$var1.pred - sqrt(loocv_data$var1.var), 
         x1 = loocv_data$observed, 
         y1 = loocv_data$var1.pred + sqrt(loocv_data$var1.var), 
         angle = 90, code = 3, length = 0.05, col = "gray")
  
  # Add reference line (x = y)
  abline(a = 0, b = 1, col = "red", lty = 2)  # Line x = y
  
  # Combine station_type and station_area into a single legend
  legend_labels <- c(names(station_type_colors), names(station_area_shapes))
  legend_colors <- c(station_type_colors, rep("gray", length(station_area_shapes)))
  legend_shapes <- c(rep(19, length(station_type_colors)), station_area_shapes)
  
  # Add combined legend
  legend("topleft", legend = legend_labels, col = legend_colors, pch = legend_shapes, 
         title = "Station Type & Area", cex = 0.8)
}
# Combine all plots into multiple pages
pdf(file = paste0(wout, "LOOCV_combined_plots_station_type_area_2024.pdf"), width = 15, height = 15)

# Define the number of plots per page
plots_per_page <- 15  # 3 rows x 5 columns
total_models <- length(loocv_results)
pages <- ceiling(total_models / plots_per_page)

# Loop through pages
for (page in 1:pages) {
  # Set up a grid layout for the current page
  par(mfrow = c(3, 5), mar = c(4, 4, 3, 1))  # 3 rows, 5 columns
  
  # Determine the range of models for this page
  start_index <- (page - 1) * plots_per_page + 1
  end_index <- min(page * plots_per_page, total_models)
  
  # Loop through the models for this page
  for (i in start_index:end_index) {
    loocv_data <- as.data.frame(loocv_results[[i]])
    
    # Add station_type and station_area to loocv_data
    loocv_data <- loocv_data %>%
      mutate(ref_obs = sf_model_traffic$ref_obs[match(
        paste(coords.x1, coords.x2),
        paste(st_coordinates(sf_model_traffic)[, 1], st_coordinates(sf_model_traffic)[, 2])
      )])
    
    loocv_data <- merge(
      loocv_data,
      sf_model_traffic %>% st_drop_geometry() %>% select(ref_obs, station_type, station_area),
      by = "ref_obs",
      all.x = TRUE
    )
    
    # Generate the plot for the current model
    cv_plot(loocv_data, model_name = model_names[i])
  }
}

# Close the PDF device
dev.off()
