####################################################################################
#                                                                                  #
#    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. Regrid the model to a finer resolution                                     #
#    4. Apply Universal Kriging for correction with traffic predictors             #                                           
#                                                                                  #
#################################################################################### 

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

# Set directories
wdir <- "/esarchive/scratch/ajulian/scripts/Open_Data/Downscaling_kriging/Data/"
wdirpredictor <-  "/esarchive/scratch/ajulian/scripts/Open_Data/Traffic_buffers/Results/"
wwind <- "/gpfs/scratch/bsc32/bsc032604/Rscripts/Downscaling/Results"
wout <- "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"      

sp_model <- readRDS(paste0(wdir, "sp_rawmodel_20220202_def.rds"))                             # CALIOPE model data  
sp_obs_locations <- readRDS(paste0(wdir,"sp_obs_locations.RDS"))                              # Obs stations coordinates
df_obs <- readRDS(paste0(wdir,"obs_20220202_8am.rds"))                                        # Obs data
domain_polygon <- readRDS(file.path(wdir, "domini_bcn.RDS"))                                  # Domain polygon
refined_grid <- readRDS(file.path(wdir, "def_refined_grid_sf_100.RDS"))                               # Refined grid
traffic_buffers <- readRDS(file.path(wdirpredictor, "traffic_buffers_norm_def.RDS"))      # Traffic buffers
wind_traffic_buffers <- readRDS(file.path(wwind, "traffic_buffers_wind_tot_20220202_8.RDS"))  # Wind traffic buffers

# Define the pollutant type variable
# (Only valid for sp_model, change manually the df_obs to match time)
pollutant_res <- "8am"  # Change this to "8am" or "24h_avg" as needed

# Map the pollutant type to the corresponding column name
pollutant_column <- switch(pollutant_res,
                           "1am" = "NO2_1am",
                           "8am" = "NO2_8am",
                           "24h_avg" = "NO2_24h_avg")

sf_model <- st_as_sf(sp_model)  # Ensure it's an sf object
sp_model@data$pollutant <- sp_model@data[[pollutant_column]]  # Rename pollutant column in sp_model

# Ensure CRS is consistent across spatial objects
sp_model <- spTransform(sp_model, CRS(st_crs(crs_utm)$proj4string))
sf_model <- st_transform(sf_model, st_crs(sp_model))
refined_grid <- st_transform(refined_grid, st_crs(sp_model))
domain_polygon <- st_transform(domain_polygon, st_crs(sp_model))

# Convert sp_obs_locations to sf, transform, and then convert to sp if needed
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(sp_model))
sp_obs_loc_bcn <- st_intersection(sp_obs_locations_sf, domain_polygon)


############################################################
#                                                          #
#     Step 2: Interpolate raw model to AQMS's locations    #
#                                                          #
############################################################

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

# Visualize the resulting variogram
plot(pollutant.fit_ok)

# Create gstat object for kriging
ok_gstat <- gstat(formula = pollutant ~ 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_loc_bcn)

# Interpolation using Ordinary Kriging at the finergrid

OK_finergrid <- predict(ok_gstat, newdata = refined_grid)

# Visualize the raw model values at the AQMS' locations
mapview(OK_obs, zcol="var1.pred")  # corrected predicted value
#mapview(OK, zcol="var1.var")            # variance of each value

# Save the interpolated model data
# saveRDS(OK_finergrid, file = paste0(wout, "OK_finergrid_20220101.RDS"))
# saveRDS(OK_obs, file = paste0(wout, "OK_obs_20220101.RDS"))

############################################################
#                                                          #
#      Step 3: Ordinary Kriging on traffic buffers         #
#                (Could not be necessary)                  #
#                                                          #
############################################################

# Ensure CRS consistency for traffic buffers
traffic_buffers <- st_transform(traffic_buffers, st_crs(sp_model))
wind_traffic_buffers <- st_transform(wind_traffic_buffers, st_crs(sp_model))

sp_traffic_buffers <- as(traffic_buffers, "Spatial")
sp_wind_traffic_buffers <- as(wind_traffic_buffers, "Spatial")
# Variogram fitting for traffic buffers
traffic100.fit_ok <- automap::autofitVariogram(formula = buffer_100m_normalized ~ 1, input_data = sp_traffic_buffers)
traffic500.fit_ok <- automap::autofitVariogram(formula = buffer_500m_normalized ~ 1, input_data = sp_traffic_buffers)
trafficwind100.fit_ok <- automap::autofitVariogram(formula = buffer_100m_normalized ~ 1, input_data = sp_wind_traffic_buffers)
trafficwind500.fit_ok <- automap::autofitVariogram(formula = buffer_500m_normalized ~ 1, input_data = sp_wind_traffic_buffers)

# Visualize the resulting variogram (Corrected variable name)
plot(traffic100.fit_ok)

# Create gstat object for traffic buffers
ok_gstat_traffic100 <- gstat(formula = buffer_100m_normalized ~ 1, data = sp_traffic_buffers, model = traffic100.fit_ok$var_model, nmax = 25)
ok_gstat_traffic500 <- gstat(formula = buffer_500m_normalized ~ 1, data = sp_traffic_buffers, model = traffic500.fit_ok$var_model, nmax = 25)
ok_gstat_trafficwind100 <- gstat(formula = buffer_100m_normalized ~ 1, data = sp_wind_traffic_buffers, model = trafficwind100.fit_ok$var_model, nmax = 25)
ok_gstat_trafficwind500 <- gstat(formula = buffer_500m_normalized ~ 1, data = sp_wind_traffic_buffers, model = trafficwind500.fit_ok$var_model, nmax = 25)

# Interpolate traffic buffers using Kriging
OK_traffic_buffers100 <- predict(ok_gstat_traffic100, newdata = OK_obs)
OK_traffic_buffers500 <- predict(ok_gstat_traffic500, newdata = OK_obs)
OK_traffic_bufferswind100 <- predict(ok_gstat_trafficwind100, newdata = OK_obs)
OK_traffic_bufferswind500 <- predict(ok_gstat_trafficwind500, newdata = OK_obs)

# Visualize the traffic buffers
mapview(OK_traffic_bufferswind500, zcol = "var1.pred")

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

# Merge OK_finergrid and traffic_buffers using st_join (spatial operation)
combined_data <- st_intersection(
  OK_finergrid, traffic_buffers)

# Rename columns in wind_traffic_buffers to avoid conflicts
colnames(wind_traffic_buffers)[colnames(wind_traffic_buffers) != "geom"] <- 
    paste0("wind_", colnames(wind_traffic_buffers)[colnames(wind_traffic_buffers) != "geom"])

# Use st_join instead of cbind to merge while preserving spatial attributes
combined_data <- st_join(combined_data, wind_traffic_buffers)

# Ensure correct column names in combined_data
colnames(combined_data)[1] <- "pollutant_model_pred"
colnames(combined_data)[2] <- "pollutant_model_var"

# Ensure all spatial datasets are using the same CRS
sf_obs_loc_bcn <- st_transform(sp_obs_loc_bcn, st_crs(domain_polygon))


# Adjusting observed data for Barcelona domain
df_rawmodel_interpolated_bcn <- data.frame(
  lat = st_coordinates(sf_obs_loc_bcn)[, 2],
  lon = st_coordinates(sf_obs_loc_bcn)[, 1],
  pollutant_model_pred = OK_obs$var1.pred,
  ref_obs = sf_obs_loc_bcn$code,
  buffer_100m_normalized = OK_traffic_buffers100$var1.pred,
  buffer_500m_normalized = OK_traffic_buffers500$var1.pred,
  wind_buffer_100m_normalized = OK_traffic_bufferswind100$var1.pred,
  wind_buffer_500m_normalized = OK_traffic_bufferswind500$var1.pred
)

# Merge with observed data
df_obs_cols <- data.frame(pollutant_obs = as.numeric(df_obs), ref_obs = colnames(df_obs))
df_model_obs_bcn <- merge(df_rawmodel_interpolated_bcn, df_obs_cols, by = "ref_obs")

# Remove missing values
df_model_obs_bcn <- na.omit(df_model_obs_bcn)

# Convert to sf object
sf_model_obs_bcn <- st_as_sf(df_model_obs_bcn, coords = c("lon", "lat"), crs = st_crs(domain_polygon))
sp_model_obs_bcn <- as(sf_model_obs_bcn, "Spatial")
# Visualize the model in Barcelona domain
mapview(sf_model_obs_bcn, zcol = "pollutant_obs")

############################################################
#                                                          #
# Step 5: Universal Kriging to the refined grid            #
#                                                          #
############################################################

# 5.1: Fit variogram for pollutant model
model.fit_uk <- automap::autofitVariogram(
  formula = pollutant_obs ~ pollutant_model_pred, 
  input_data = sp_model_obs_bcn,
  verbose = TRUE,
  miscFitOptions = list(merge.small.bins = FALSE)
)
uk_gstat_model <- gstat(
  formula = pollutant_obs ~ pollutant_model_pred, 
  data = sp_model_obs_bcn,
  model = model.fit_uk$var_model
)

refined_corrected_model <- predict(object = uk_gstat_model, newdata = combined_data)

colnames(refined_corrected_model)[1] <- "UK_model_pred"
colnames(refined_corrected_model)[2] <- "UK_model_var"

# Visualize the corrected model
mapview(refined_corrected_model, zcol = "UK_model_pred")
mapview(refined_corrected_model, zcol = "UK_model_var")

# 5.2: Fit variogram for traffic buffers
traffic100.fit_uk <- automap::autofitVariogram(
  formula = pollutant_obs ~ pollutant_model_pred + buffer_100m_normalized, 
  input_data = sp_model_obs_bcn,
  verbose = TRUE,
  miscFitOptions = list(merge.small.bins = FALSE)
)
uk_gstat_traffic100 <- gstat(
  formula = pollutant_obs ~ pollutant_model_pred + buffer_100m_normalized, 
  data = sp_model_obs_bcn,
  model = traffic100.fit_uk$var_model
)

refined_corrected_traffic100 <- predict(object = uk_gstat_traffic100, newdata = combined_data)


traffic500.fit_uk <- automap::autofitVariogram(
  formula = pollutant_obs ~ pollutant_model_pred + buffer_500m_normalized, 
  input_data = sp_model_obs_bcn,
  verbose = TRUE,
  miscFitOptions = list(merge.small.bins = FALSE)
)
uk_gstat_traffic500 <- gstat(
  formula = pollutant_obs ~ pollutant_model_pred + buffer_500m_normalized, 
  data = sp_model_obs_bcn,
  model = traffic500.fit_uk$var_model
)

refined_corrected_traffic500 <- predict(object = uk_gstat_traffic500, newdata = combined_data)

mapview(refined_corrected_traffic500, zcol = "var1.pred")
# 5.3: Fit variogram for wind traffic buffers

trafficwind100.fit_uk <- automap::autofitVariogram(
  formula = pollutant_obs ~ pollutant_model_pred + wind_buffer_100m_normalized, 
  input_data = sp_model_obs_bcn,
  verbose = TRUE,
  miscFitOptions = list(merge.small.bins = FALSE)
)
uk_gstat_windtraffic100 <- gstat(
  formula = pollutant_obs ~ pollutant_model_pred + wind_buffer_100m_normalized, 
  data = sp_model_obs_bcn,
  model = trafficwind100.fit_uk$var_model
)

refined_corrected_windtraffic100 <- predict(object = uk_gstat_windtraffic100, newdata = combined_data)

trafficwind500.fit_uk <- automap::autofitVariogram(
  formula = pollutant_obs ~ pollutant_model_pred + wind_buffer_500m_normalized, 
  input_data = sp_model_obs_bcn,
  verbose = TRUE,
  miscFitOptions = list(merge.small.bins = FALSE)
)
uk_gstat_windtraffic500 <- gstat(
  formula = pollutant_obs ~ pollutant_model_pred + wind_buffer_500m_normalized, 
  data = sp_model_obs_bcn,
  model = trafficwind500.fit_uk$var_model
)

refined_corrected_windtraffic500 <- predict(object = uk_gstat_windtraffic500, newdata = combined_data)

combined_refined_corrected <- refined_corrected_model 
combined_refined_corrected$UK_traffic100_pred <- refined_corrected_traffic100$var1.pred
combined_refined_corrected$UK_traffic500_pred <- refined_corrected_traffic500$var1.pred
combined_refined_corrected$UK_windtraffic100_pred <- refined_corrected_windtraffic100$var1.pred
combined_refined_corrected$UK_windtraffic500_pred <- refined_corrected_windtraffic500$var1.pred
combined_refined_corrected$UK_traffic100_var <- refined_corrected_traffic100$var1.var
combined_refined_corrected$UK_traffic500_var <- refined_corrected_traffic500$var1.var
combined_refined_corrected$UK_windtraffic100_var <- refined_corrected_windtraffic100$var1.var
combined_refined_corrected$UK_windtraffic500_var <- refined_corrected_windtraffic500$var1.var

# Visualize the corrected model
mapview(combined_refined_corrected, zcol = "UK_windtraffic100_pred")

# 5.4: Fit variogram for combination traffic buffers

trafficcomb.fit_uk <- automap::autofitVariogram(
  formula = pollutant_obs ~ pollutant_model_pred + buffer_100m_normalized + buffer_500m_normalized, 
  input_data = sp_model_obs_bcn,
  verbose = TRUE,
  miscFitOptions = list(merge.small.bins = FALSE)
)
uk_gstat_traffic_comb <- gstat(
  formula = pollutant_obs ~ pollutant_model_pred + wind_buffer_100m_normalized + buffer_500m_normalized, 
  data = sp_model_obs_bcn,
  model = trafficcomb.fit_uk$var_model
)

refined_corrected_traffic_comb <- predict(object = uk_gstat_traffic_comb, newdata = combined_data)
mapview(refined_corrected_traffic_comb, zcol = "var1.pred")

trafficwindcomb.fit_uk <- automap::autofitVariogram(
  formula = pollutant_obs ~ pollutant_model_pred + wind_buffer_100m_normalized + wind_buffer_500m_normalized, 
  input_data = sp_model_obs_bcn,
  verbose = TRUE,
  miscFitOptions = list(merge.small.bins = FALSE)
)
uk_gstat_wind_traffic_comb <- gstat(
  formula = pollutant_obs ~ pollutant_model_pred + wind_buffer_100m_normalized + wind_buffer_500m_normalized, 
  data = sp_model_obs_bcn,
  model = trafficwindcomb.fit_uk$var_model
)

############################################################
#                                                          #
#      Step 6: leave-one-out cross validation              #
#                                                          #
############################################################

loocvmodel.fit_ok <- krige.cv(
  pollutant_model_pred ~ 1, 
  sf_model_obs_bcn, 
  pollutant.fit_ok$var_model, 
  nmax = 25)

loocvmodel.fit_uk <- krige.cv(
  pollutant_obs ~ pollutant_model_pred, 
  sf_model_obs_bcn,
  model.fit_uk$var_model
  )

loocvtraffic100.fit_uk <- krige.cv(
    pollutant_obs ~ pollutant_model_pred + buffer_100m_normalized, 
    sf_model_obs_bcn,
    traffic100.fit_uk$var_model
    )

loocvtraffic500.fit_uk <- krige.cv(
    pollutant_obs ~ pollutant_model_pred + buffer_500m_normalized, 
    sf_model_obs_bcn,
    traffic500.fit_uk$var_model
    )

loocvwindtraffic100.fit_uk <- krige.cv(
    pollutant_obs ~ pollutant_model_pred + wind_buffer_100m_normalized, 
    sf_model_obs_bcn,
    trafficwind100.fit_uk$var_model
    )

loocvwindtraffic500.fit_uk <- krige.cv(
    pollutant_obs ~ pollutant_model_pred + wind_buffer_500m_normalized, 
    sf_model_obs_bcn,
    trafficwind500.fit_uk$var_model
    )

loocvtrafficcomb.fit_uk <- krige.cv(
    pollutant_obs ~ pollutant_model_pred + buffer_100m_normalized + buffer_500m_normalized, 
    sf_model_obs_bcn,
    trafficcomb.fit_uk$var_model
    )

loocvwindtrafficcomb.fit_uk <- krige.cv(
    pollutant_obs ~ pollutant_model_pred + wind_buffer_100m_normalized + wind_buffer_500m_normalized, 
    sf_model_obs_bcn,
    trafficwindcomb.fit_uk$var_model
    )

# List of LOOCV results
loocv_results <- list(
  loocvmodel.fit_ok, 
  loocvmodel.fit_uk, 
  loocvtraffic100.fit_uk, 
  loocvtraffic500.fit_uk,
  loocvwindtraffic100.fit_uk,
  loocvwindtraffic500.fit_uk,
  loocvtrafficcomb.fit_uk,
  loocvwindtrafficcomb.fit_uk
)

# Model names for easy identification
model_names <- c(
  "fit_ok",
  "fit_uk",
  "traffic100_uk",
  "traffic500_uk",
  "windtraffic100_uk",
  "windtraffic500_uk",
  "trafficcomb_uk",
  "windtrafficcomb_uk"
)

# Initialize an empty data frame to store results
stats_results <- data.frame()

# Loop through each LOOCV model, extract observed and predicted values, and calculate stats
for (i in 1:length(loocv_results)) {
  
  # Extract observed and predicted values
  loocv_data <- data.frame(
    obs = loocv_results[[i]]$observed, # observed values
    pred = loocv_results[[i]]$var1.pred # predicted values
  )
  
  # Calculate the statistics using modStats
  stats <- modStats(loocv_data, mod = "pred", obs = "obs")
  
  # Add a column for the model name
  stats$model <- model_names[i]
  
  # Append the result to the stats_results data frame
  stats_results <- rbind(stats_results, stats)
}

# Display the result
print(stats_results)

