# =============================================================================
# Script Name: Datafusion.R
# Description: This script performs spatial interpolation, cross-validation, 
#              and regridding for NO2 concentration data using Ordinary Kriging (OK) 
#              and Universal Kriging (UK). It includes variogram fitting, 
#              interpolation, regridding, and model evaluation.
# 
# Author: Andreu Julian Izquierdo
# Date Created: 28/03/2025
# Last Updated: 08/04/2025
# 
# Input:
#   - RDS files containing NO2 model data ("mean_campaign_c*.RDS")
#   - RDS files with observational data ("sp_coord_dosimeters.rds", "df_obs_dosimeters.rds")
# 
# Output:
#   - Interpolated NO2 concentration maps
#   - Cross-validation statistics for OK and UK
#   - Variograms for raw model, observations, and corrected model
#   - Regridded rasters for mean, standard deviation, and relative standard deviation
#   - Combined visualization of regridded rasters
# 
# Steps:
#   1. Load NO2 model and observational data.
#   2. Perform Ordinary Kriging (OK) to interpolate NO2 concentrations.
#   3. Perform Universal Kriging (UK) to correct model predictions using observations.
#   4. Evaluate the models using cross-validation.
#   5. Perform cross-validation for Universal Kriging.
#   6. Analyze variograms for raw model, observations, and corrected model.
#   7. Regrid data to a uniform grid using Ordinary Kriging and plot the results.
# =============================================================================

# Load libraries
library(sf)
library(sp)
library(gstat)
library(automap)
library(mapview)
library(openair)
library(ggplot2)
#library(ggpubr)
#library(lattice)
library(raster)
library(cowplot)
library(rasterVis)  # For raster plotting with ggplot2

#library(osmdata)



# Setting
wdir <- "/esarchive/scratch/ajulian/Aveiro/Data/5m_nmd_mean/Period_5/"
#wdir <- "/Users/ajulian/Documents/Aveiro/CALIOPE_Urban_sim/5m/Period_2/"
wdir2 <- "/esarchive/scratch/ajulian/Aveiro/Scripts/"
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 obs and model                           #
#                                                          #
############################################################
sp_model <-readRDS(paste0(wdir,"mean_campaign_c5.rds"))
sp_obs_locations <- readRDS(paste0(wdir2,"sp_coord_dosimeters_c5.rds"))
df_obs <- readRDS(paste0(wdir2,"df_obs_dosimeters_c5.rds"))
buildings <- readRDS(paste0(wdir2, "buildings_cropped_closer.rds"))
bbox <- readRDS(paste0(wdir2, "expanded_bbox.rds"))

#mapview(sp_model, zcol="NO2_mean")

sp_model <- st_transform(sp_model, crs_utm)
buildings <- st_transform(buildings, crs_utm)
bbox <- st_transform(bbox, crs_utm)

sp_model <- st_intersection(sp_model, bbox)
sp_model <- as(sp_model,"Spatial")
############################################################
#                                                          #
#     Step 2: interpolate raw to AQMS's location           #
#                                                          #
############################################################

#Compute Variogram
#we are picking random points to ease the variogram computation (there are >8000 points)
#set.seed(1)
#pick_random_points        <- sample(1:nrow(sp_model),1000)
# Variogram fitting
pollutant.fit_ok <- automap::autofitVariogram(
  formula = NO2_mean ~ 1,
  input_data = sp_model
)
# Visualize the resulting variogram
plot(pollutant.fit_ok)

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

# project the spatial points dataframe to indicate where to interpolate in UTM
sp_obs <- st_transform(sp_obs_locations, st_crs(sp_model))

# Interpolation using Ordinary Kriging
OK   <- predict(object=ok_gstat, newdata=sp_obs)

# Visualize the raw model values at the AQMS' locations
mapview(OK,zcol="var1.pred")

############################################################
#                                                          #
#    Step 3&4: Universal Kriging                           #
#                                                          #
############################################################

# Build a dataframe with model&obs information
df_rawmodel_interpolated <- data.frame(
  lat = st_coordinates(sp_obs)[, 2],  # Extract latitude
  lon = st_coordinates(sp_obs)[, 1],  # Extract longitude
  pollutant_model = OK$var1.pred,     # Access the predicted values
  Site = sp_obs_locations$Site        # Access the "Site" column from the original data
)
#df_obs_cols <- data.frame(pollutant_obs = as.numeric(df_obs[1,]) , Site = colnames(df_obs))
df_model_obs <- merge(df_rawmodel_interpolated, df_obs, by = "Site")

# Visualize model & obs
head(df_model_obs)


# Plot the observed vs. raw model values
plot(df_model_obs$pollutant_obs, df_model_obs$pollutant_model, 
     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))

# Add the line x = y
abline(a = 0, b = 1, col = "red", lty = 2)  # Red dashed line

# Close the graphics device
dev.off()

# Optionally, print a message to confirm the save
print("Observed vs. Raw Model plot has been saved as 'observed_vs_raw_model_plot.png'")
# Remove missing values
df_model_obs <- na.omit(df_model_obs)
#View(df_model_obs)

# Build spatial points dataframe
sp_model_obs <- df_model_obs
coordinates(sp_model_obs)        <- ~lon+lat
sp::proj4string(sp_model_obs)    <- st_crs(sp_model)$proj4string
#mapview(sp_model_obs)

colnames(sp_model_obs@data)[colnames(sp_model_obs@data) == "Results (ug/m3)"] <- "pollutant_obs"


# Fit variogram of the residual regresision model between obs and raw CALIOPE

pollutant.fit_loocv <- automap::autofitVariogram(formula=pollutant_obs~pollutant_model,input_data=sp_model_obs, verbose=FALSE, miscFitOptions = list(merge.small.bins = FALSE))
plot(pollutant.fit_loocv)

colnames(OK)[colnames(OK) == "var1.pred"] <- "pollutant_model"

names(sp_model)[names(sp_model) == "NO2_mean"] <- "pollutant_model"

# Apply Universal Kriging to the whole map
uk_gstat          <- gstat(formula=pollutant_obs~pollutant_model,data=sp_model_obs, model = pollutant.fit_loocv$var_model)
corrected_model  <- predict(object=uk_gstat, newdata=sp_model)

mapview(corrected_model, zcol="var1.pred")
mapview(corrected_model, zcol="var1.var")# Plot the standard deviation (sqrt(var1.var)) of the corrected model
corrected_model$std_dev <- sqrt(corrected_model$var1.var)  # Calculate the standard deviation
mapview(corrected_model, zcol = "std_dev")  # Visualize the standard deviation

# Save the corrected_model as an RDS file
#saveRDS(corrected_model, file = paste0(wdir2, "corrected_model_results.RDS"))

# Optionally, print a message to confirm the save
print("Corrected model results have been saved to 'corrected_model_results.RDS'")

############################################################
#                                                          #
#     Step 5: Cross Validation                             #
#                                                          #
############################################################

# Cross-validation for Universal Kriging
cv_uk <- krige.cv(pollutant_obs~pollutant_model, sp_model_obs, model=pollutant.fit_loocv$var_model)

# Plot the observed vs. predicted values for Universal Kriging
# Plot the observed vs. predicted values for Universal Kriging
cv_plot <- function() {
  plot(cv_uk$observed, cv_uk$var1.pred, 
       xlab = expression("Observed (" * mu * "g/m"^3 * ")"),  # Add units to the x-axis label
       ylab = expression("Predicted (" * mu * "g/m"^3 * ")"),  # Add units to the y-axis label
       main = "Cross-validation: Universal Kriging", 
       xlim = c(0, max(cv_uk$observed, cv_uk$var1.pred, na.rm = TRUE)), 
       ylim = c(0, max(cv_uk$observed, cv_uk$var1.pred, na.rm = TRUE)))
  
  # Add vertical error bars
  arrows(x0 = cv_uk$observed, 
         y0 = cv_uk$var1.pred - sqrt(cv_uk$var1.var), 
         x1 = cv_uk$observed, 
         y1 = cv_uk$var1.pred + sqrt(cv_uk$var1.var), 
         angle = 90, code = 3, length = 0.05, col = "blue")
  
  # Add the line x = y
  abline(a = 0, b = 1, col = "red", lty = 2)  # Red dashed line
  # Add dashed lines for ±25% of the x = y line
  abline(a = 0, b = 1.25, col = "darkgreen", lty = 2)  # +25% line
  abline(a = 0, b = 0.75, col = "darkgreen", lty = 2)  # -25% line
  
  # Add dashed lines for ±1.5% of the x = y line
  #abline(a = 0, b = 1.015, col = "orange", lty = 2)  # +1.5% line
  #abline(a = 0, b = 0.985, col = "orange", lty = 2)  # -1.5% line
}

# Save the plot as a PNG file
#png(filename = paste0(wdir2, "cv_uk_plot.png"), width = 800, height = 800, res = 150)
cv_plot()
dev.off()


# Optionally, print a message to confirm the save
print("Cross-validation plot has been saved as 'cv_uk_plot.png'")
cv_uk_df <- data.frame(
  observed = cv_uk$observed,
  predicted = cv_uk$var1.pred
)
uk_stats <- modStats(cv_uk_df, mod="predicted", obs="observed")
print("Model Statistics for Universal Kriging:")
print(uk_stats)


############################################################
#                                                          #
#     Step 6: Variogram of Observations and Corrected Model#
#                                                          #
############################################################

# Variogram for raw model (pollutant.fit_ok)
plot(pollutant.fit_ok)
print("Variogram for Raw Model:")
print(pollutant.fit_ok$var_model)


# Variogram for observational data (df_obs)
obs_variogram <- automap::autofitVariogram(
  formula = pollutant_obs ~ 1, 
  input_data = sp_model_obs,
  nugget = 0.25*(mean(sp_model_obs$pollutant_obs)))

plot(obs_variogram)

print("Variogram for Observational Data:")
print(obs_variogram$var_model)

# Variogram for corrected model (Universal Kriging output)
corrected_variogram <- automap::autofitVariogram(
  formula = var1.pred ~ 1, 
  input_data = corrected_model)
plot(corrected_variogram)
print("Variogram for Corrected Model:")
print(corrected_variogram$var_model)

############################################################
#                                                          #
#     Step 7: Regrid Data Using Ordinary Kriging           #
#                                                          #
############################################################

regrid_ok <- function(non_uniform_data, target_grid, crs_sim = "+proj=utm +zone=31 +ellps=intl +units=m +no_defs") {
  # Ensure the input data is a SpatialPointsDataFrame
  if (!inherits(non_uniform_data, "SpatialPointsDataFrame")) {
    non_uniform_data <- sp::SpatialPointsDataFrame(non_uniform_data[, c("X", "Y")], non_uniform_data)
    sp::proj4string(non_uniform_data) <- crs_sim
  }
  
  # Fit the variogram for Ordinary Kriging
  vf_ok <- automap::autofitVariogram(pollutant_uk ~ 1, non_uniform_data)
  
  # Perform Ordinary Kriging
  ok_regular <- gstat(formula = pollutant_uk ~ 1, data = non_uniform_data, model = vf_ok$var_model, nmax = 20)
  regular <- predict(ok_regular, target_grid)
  
  # Convert the result to a raster
  mean_raster <- raster::rasterFromXYZ(as.data.frame(regular)[, c("x", "y", "var1.pred")])
  crs(mean_raster) <- crs_sim  # Assign CRS to the raster
  
  return(mean_raster)
}

# Function to regrid standard deviation
regrid_ok_sd <- function(non_uniform_data, target_grid, crs_sim = "+proj=utm +zone=31 +ellps=intl +units=m +no_defs") {
  # Remove rows with NA in the sd column
  non_uniform_data <- subset(non_uniform_data, !is.na(non_uniform_data$sd))
  
  # Ensure the input data is a SpatialPointsDataFrame
  if (!inherits(non_uniform_data, "SpatialPointsDataFrame")) {
    non_uniform_data <- sp::SpatialPointsDataFrame(non_uniform_data[, c("X", "Y")], non_uniform_data)
    sp::proj4string(non_uniform_data) <- crs_sim
  }
  
  # Fit the variogram for Ordinary Kriging
  vf_ok <- automap::autofitVariogram(sd ~ 1, non_uniform_data)
  
  # Perform Ordinary Kriging
  ok_regular <- gstat(formula = sd ~ 1, data = non_uniform_data, model = vf_ok$var_model, nmax = 20)
  regular <- predict(ok_regular, target_grid)
  
  # Convert the result to a raster
  mean_raster <- raster::rasterFromXYZ(as.data.frame(regular)[, c("x", "y", "var1.pred")])
  crs(mean_raster) <- crs_sim  # Assign CRS to the raster
  
  return(mean_raster)
}

# Function to regrid relative standard deviation
regrid_ok_sd_rel <- function(non_uniform_data, target_grid, crs_sim = "+proj=utm +zone=31 +ellps=intl +units=m +no_defs") {
  # Remove rows with NA in the sd_rel column
  non_uniform_data <- subset(non_uniform_data, !is.na(non_uniform_data$sd_rel))
  
  # Ensure the input data is a SpatialPointsDataFrame
  if (!inherits(non_uniform_data, "SpatialPointsDataFrame")) {
    non_uniform_data <- sp::SpatialPointsDataFrame(non_uniform_data[, c("X", "Y")], non_uniform_data)
    sp::proj4string(non_uniform_data) <- crs_sim
  }
  
  # Fit the variogram for Ordinary Kriging
  vf_ok <- automap::autofitVariogram(sd_rel ~ 1, non_uniform_data)
  
  # Perform Ordinary Kriging
  ok_regular <- gstat(formula = sd_rel ~ 1, data = non_uniform_data, model = vf_ok$var_model, nmax = 20)
  regular <- predict(ok_regular, target_grid)
  
  # Convert the result to a raster
  mean_raster <- raster::rasterFromXYZ(as.data.frame(regular)[, c("x", "y", "var1.pred")])
  crs(mean_raster) <- crs_sim  # Assign CRS to the raster
  
  return(mean_raster)
}

corrected_model_sp <- as(corrected_model, "Spatial")  # Convert to SpatialPixelsDataFrame

sp_model_sp <- as(sp_model, "Spatial")

# Ara pots fer servir la funció regrid_ok
mean_raster_nc <- regrid_ok(sp_model_sp, target_grid_sp)
# Define a target grid for regridding
target_grid <- raster::raster(raster::extent(corrected_model_sp), res = 1)  # 20m resolution
raster::crs(target_grid) <- raster::crs(corrected_model_sp)  # Set CRS to match the data

# Add a column for pollutant_uk (predicted values) to corrected_model_sp
corrected_model_sp@data$pollutant_uk <- corrected_model_sp$var1.pred

# Convert the target grid to a SpatialPixelsDataFrame
target_grid_sp <- as(target_grid, "SpatialPixelsDataFrame")

# Regrid the mean, standard deviation, and relative standard deviation
mean_raster_regridded <- regrid_ok(corrected_model_sp, target_grid_sp)

# Assuming `corrected_model_sp` contains `sd` and `sd_rel` columns
corrected_model_sp@data$sd <- corrected_model_sp$std_dev  # Assign standard deviation
corrected_model_sp@data$sd_rel <- corrected_model_sp$std_dev / corrected_model_sp$var1.pred  # Calculate relative standard deviation
sp_model_sp@data$pollutant_uk <- sp_model_sp@data$pollutant_model

# Regrid standard deviation
sd_raster_regridded <- regrid_ok_sd(corrected_model_sp, target_grid_sp)

# Regrid relative standard deviation
sd_rel_raster_regridded <- regrid_ok_sd_rel(corrected_model_sp, target_grid_sp)

mean_raster_nc <- regrid_ok(sp_model_sp, target_grid_sp)

# Convert rasters to data frames for ggplot2
mean_df <- as.data.frame(mean_raster_regridded, xy = TRUE)
colnames(mean_df)[3] <- "layer"  # Rename the raster value column to "layer"

sd_df <- as.data.frame(sd_raster_regridded, xy = TRUE)
colnames(sd_df)[3] <- "layer"  # Rename the raster value column to "layer"

sd_rel_df <- as.data.frame(sd_rel_raster_regridded, xy = TRUE)
colnames(sd_rel_df)[3] <- "layer"  # Rename the raster value column to "layer"

mean_df_nc <- as.data.frame(mean_raster_nc, xy = TRUE)
colnames(mean_df_nc)[3] <- "layer"

# Ensure CRS of sp_obs_locations matches the raster
sp_obs_locations <- st_transform(sp_obs_locations, crs = raster::crs(mean_raster_regridded))

# Convert sp_obs_locations to a data frame for ggplot2
sp_obs_locations_df <- as.data.frame(st_coordinates(sp_obs_locations))
colnames(sp_obs_locations_df) <- c("x", "y")

# Create the plot
mean_plot <- ggplot() +
  geom_raster(data = mean_df, aes(x = x, y = y, fill = layer)) +
  geom_sf(data = buildings, fill = NA, color = "black", size = 0.2) +  # Overlay buildings
  geom_sf(data = sp_obs_locations, color = "grey", size = 0.5) +  # Overlay observation locations with fixed color
  scale_fill_viridis_c(
    name = expression(mu * g/m^3),
    option = "C",
    guide = guide_colorbar(title.position = "top", title.hjust = 0.5)  # Center the title above the color bar
  ) +
  coord_sf(crs = st_crs(buildings)) +  # Use coord_sf() to handle spatial data
  theme_minimal(base_size = 11) +
  theme(
    plot.title = element_text(hjust = 0.5, size = 14, face = "bold"),  # Centered title
    axis.title = element_blank(),
    axis.text = element_blank(),
    axis.ticks = element_blank(),
    legend.position = "right",  # Move the legend to the right
    legend.text = element_text(size = 8),
    legend.title = element_text(size = 9)
  ) 


# Print the plot
print(mean_plot)
#ggsave("corrected_NO2_with_buildings_c1.png", mean_plot, width = 6, height = 6, dpi = 800)



sd_plot <- ggplot() +
  geom_raster(data = sd_df, aes(x = x, y = y, fill = layer)) +
  geom_sf(data = buildings, fill = NA, color = "black", size = 0.2) +  # Overlay buildings
  geom_sf(data = sp_obs_locations, color = "grey", size = 0.5) +  # Overlay observation locations
  scale_fill_viridis_c(
    name = expression(mu * g/m^3),
    option = "C",
    guide = guide_colorbar(title.position = "top", title.hjust = 0.5)  # Center the title above the color bar
  ) +
  coord_sf(crs = st_crs(buildings)) +  # Use coord_sf() to handle spatial data
  theme_minimal(base_size = 11) +
  theme(
    plot.title = element_text(hjust = 0.5, size = 14, face = "bold"),  # Centered title
    axis.title = element_blank(),
    axis.text = element_blank(),
    axis.ticks = element_blank(),
    legend.position = "right",  # Move the legend to the right
    legend.text = element_text(size = 8),
    legend.title = element_text(size = 9)
  ) 
print(sd_plot)
#ggsave("sd_NO2_with_buildings_c1.png", sd_plot, width = 6, height = 6, dpi = 800)

sd_rel_plot <- ggplot() +
  geom_raster(data = sd_rel_df, aes(x = x, y = y, fill = layer)) +
  geom_sf(data = buildings, fill = NA, color = "black", size = 0.2) +  # Overlay buildings
  geom_sf(data = sp_obs_locations, color = "grey", size = 0.5) +  # Overlay observation locations
  scale_fill_viridis_c(
    name = "Relative SD",
    option = "C",
    guide = guide_colorbar(title.position = "top", title.hjust = 0.5)  # Center the title above the color bar
  ) +
  coord_sf(crs = st_crs(buildings)) +  # Use coord_sf() to handle spatial data
  theme_minimal(base_size = 11) +
  theme(
    plot.title = element_text(hjust = 0.5, size = 14, face = "bold"),  # Centered title
    axis.title = element_blank(),
    axis.text = element_blank(),
    axis.ticks = element_blank(),
    legend.position = "right",  # Move the legend to the right
    legend.text = element_text(size = 8),
    legend.title = element_text(size = 9)
  )
print(sd_rel_plot)
#ggsave("sd_rel_NO2_with_buildings_c1.png", sd_rel_plot, width = 6, height = 6, dpi = 800)


# Combine the plots into a single row with spacing
combined_plot <- plot_grid(
  mean_plot, sd_plot, sd_rel_plot,
  nrow = 1,
  labels = c("Corrected model", "Standard deviation", "Relative standard deviation"),  # Add labels to the plots
  label_size = 8,
  rel_widths = c(1, 1, 1)  # Equal widths for all plots
)
# Print the combined plot
print(combined_plot)

mean_plot_nc <- ggplot() +
  geom_raster(data = mean_df_nc, aes(x = x, y = y, fill = layer)) +
  geom_sf(data = buildings, fill = NA, color = "black", size = 0.2) +
  geom_sf(data = sp_obs_locations, color = "grey", size = 0.5) +
  scale_fill_viridis_c(name = expression(mu * g/m^3), option = "C",
                       guide = guide_colorbar(title.position = "top", title.hjust = 0.5)) +
  coord_sf(crs = st_crs(buildings)) +
  theme_minimal(base_size = 11) +
  theme(plot.title = element_text(hjust = 0.5, size = 14, face = "bold"),
        axis.title = element_blank(), axis.text = element_blank(), axis.ticks = element_blank(),
        legend.position = "right", legend.text = element_text(size = 8), legend.title = element_text(size = 9))


print(mean_plot_nc)




# Save the combined plot as a PNG file with a horizontal layout
#ggsave("combined_rasters_horizontal_c1.png", combined_plot, width = 18, height = 6, dpi = 600)


# Plots


