####################################################################################
#                                                                                  #
#    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 libraries
library(gstat)
library(sp)
library(sf)
library(automap)
library(viridis)
library(akima)
library(mapview)



# Setting
wdir <- "/esarchive/scratch/ajulian/scripts/Open_Data/Downscaling_kriging/Data/"
crs_latlon <- "+proj=longlat +datum=WGS84 +ellps=WGS84 +towgs84=0,0,0" # projection data
crs_utm    <- "+proj=utm +zone=31 +ellps=intl +units=m +no_defs"       # universal transversal mercator zona 31 cat

############################################################
#                                                          #
#     Step 1: load obs and model                           #
#                                                          #
############################################################

sp_model <-readRDS(paste0(wdir,"sp_model.RDS"))                                   # load CALIOPE model data
sp_obs_locations <- readRDS(paste0(wdir,"sp_obs_locations.RDS"))                  # load obs stations coordinates
df_obs <- readRDS(paste0(wdir,"df_obs.RDS"))                                      # load obs data 
df_obs_filled <- readRDS(paste0(wdir,"df_obs_filled.RDS"))                        # load gap-filling data 
domain_polygon <- readRDS(file.path(wdir, "domini_bcn.RDS"))             # load domain polygon
refined_grid <- readRDS(file.path(wdir, "refined_grid_sf.RDS"))          # load refined grid
traffic_buffers <- readRDS(file.path(wdir, "traffic_buffers_bcn.RDS"))   # load traffic buffers



# in case we want to run the code with the filled data
#df_obs <- df_obs_filled 

############################################################
#                                                          #
#     Step 2: interpolate raw to AQMS's location           #
#                                                          #
############################################################

# compute Variogram
# we are picking random points to ease the variogram computation (there are >8000 points)  (mostra dels 8000 punts, n'agafem 1000)
set.seed(1)                 # to make it reproducible and give the same results

pick_random_points        <- sample(1:nrow(sp_model),1000)  # create a vector with 1000 points with random values from 1 to the total number of points in the model

input_data_example <- sp_model[pick_random_points, ] # select the points we had in the model randomly generated (specific rows, all columns)

#mapview(input_data_example) # show the points

# Variogram fitting (weights of each station)
pollutant.fit_ok          <- automap::autofitVariogram(formula=pollutant_model~1, input_data=sp_model[pick_random_points,],model="Gau")  #agafem 1000 punts per fer el variograma (si els agafessim tots seria molt pesat)

# Visualize the resulting variogram
#plot(pollutant.fit_ok)  

# Set gstat object  (pollutant~1 lineal estimation)                   
ok_gstat<- gstat(formula=pollutant_model~1, data = sp_model, model = pollutant.fit_ok$var_model, nmax=25)  # transform the variogram to a gstat (geostatistics)

# project the spatial points dataframe to indicate where to interpolate in UTM
sp_obs <- spTransform(sp_obs_locations, sp::CRS(crs_utm))  

# Interpolation using Ordinary Kriging
OK   <- predict(object=ok_gstat, newdata=sp_obs)   # do the ordinary kriging from the variogram transformed to gstats, and the new data is the data of the observations in all the stations
# it returns the predicted value of the model with the correction thanks to the variogram and the variance in the observation places

# Convert OK to sf object
OK_sf <- st_as_sf(OK)
# Visualize the raw model values at the AQMS' locations 
# mapview(OK, zcol="var1.pred")  # corrected predicted value
# mapview(OK, zcol="var1.var")   # variance of each value

############################################################
#                                                          #
#    Step 2.1: Regriding model in Barcelona                #
#                                                          #
############################################################

# Ensure refined_grid and domain_polygon has the same CRS as ok_gstat
refined_grid <- st_transform(refined_grid, st_crs(sp_model))
domain_polygon <- st_transform(domain_polygon, st_crs(sp_model))

OK_sf <- st_transform(OK_sf, st_crs(domain_polygon))
OK_bcn <- st_intersection(OK_sf, domain_polygon)

# Interpolation using Ordinary Kriging
OK_finergrid <- predict(object = ok_gstat, newdata = refined_grid)  # ordinary kriging to the new refined grid  

# Crop the data to the domain of interest using st_intersection
sp_model_bcn <- st_intersection(OK_finergrid, domain_polygon)

# Visualize the corrected model
mapview(OK_finergrid, zcol = "var1.pred")
#OK_finergrid is a data frame with the predicted values of the model in the refined grid

###############################################################
#                                                             #
# Step 3: Ordinary kriging on the obs for the traffic buffers #                          
#                                                             #
###############################################################

traffic_buffers <- st_transform(traffic_buffers, crs = crs_utm)

# Variogram fitting (weights of each station)
traffic.fit_ok          <- automap::autofitVariogram(formula=buffer_500m~1, input_data=traffic_buffers,model="Gau")  #agafem 1000 punts per fer el variograma (si els agafessim tots seria molt pesat)
#plot(traffic.fit_ok)
# Visualize the resulting variogram
#plot(pollutant.fit_ok)  

# Set gstat object  (pollutant~1 lineal estimation)                   
ok_gstat_traffic<- gstat(formula=buffer_500m~1, data = traffic_buffers, model = traffic.fit_ok$var_model, nmax=25)  # transform the variogram to a gstat (geostatistics)

# project the spatial points dataframe to indicate where to interpolate in UTM
sp_obs <- spTransform(sp_obs_locations, sp::CRS(crs_utm))  

#Interpolation using Ordinary Kriging
OK_traffic_buffers <- predict(object=ok_gstat_traffic, newdata=OK)

OK_finergrid_traffic_buffers <- predict(object = ok_gstat_traffic, newdata = refined_grid)  # ordinary kriging to the new refined grid
# Convert OK to sf object
OK_finergrid_traffic_buffers <- st_as_sf(OK_finergrid_traffic_buffers)

mapview(OK_finergrid_traffic_buffers, zcol = "var1.pred")
  
###########
#Merge OK_finergrid and OK_traffic_buffers
###########

combined_data <- st_intersection(OK_finergrid, OK_finergrid_traffic_buffers)

#########
# Adjusting obs data in the Barcelona domain
#########
mapview::mapview(OK_traffic_buffers, zcol="var1.pred")
# Adjusting obs data in the Barcelona domain
df_rawmodel_interpolated               <- data.frame(lat=sp_obs_locations@coords[,2],
                                           lon=sp_obs_locations@coords[,1],
                                           pollutant_model=OK@data$var1.pred,
                                           ref_obs=sp_obs_locations@data$code,
                                           buffer_500m=OK_traffic_buffers$var1.pred)
df_obs_cols <- data.frame(pollutant_obs = as.numeric(df_obs[1,]) , ref_obs = colnames(df_obs))
df_model_obs <- merge(df_rawmodel_interpolated, df_obs_cols, by = "ref_obs")

# Remove missing values
df_model_obs <- na.omit(df_model_obs) 
# Build spatial points dataframe 
sp_model_obs <- df_model_obs 

# Set the coordinates of the spatial points dataframe
coordinates(sp_model_obs)        <- ~lon+lat
sp::proj4string(sp_model_obs)    <- crs_utm

sf_model_obs <- st_as_sf(df_model_obs, coords = c("lon", "lat"), crs = crs_utm)

######
# Crop to Barcelona city
######

sf_model_obs <- st_transform(sf_model_obs, st_crs(domain_polygon))
# Crop the data to the domain of interest using st_intersection
sf_model_obs_bcn <- st_intersection(sf_model_obs, domain_polygon)

mapview(sf_model_obs_bcn, zcol = "pollutant_model")


###############################################################
#                                                             #
# Step 4: Universal Kriging based on the new refined grid     #                          
#                                                             #
###############################################################

# Build a dataframe with model&obs information

# Check the structure of the combined data
#print(combined_data)

# Fit a variogram of the residual regression model between obs and raw CALIOPE
pollutant.fit_loocv_traffic <- automap::autofitVariogram(
  formula = pollutant_obs ~ pollutant_model + buffer_500m,  # Include traffic as predictor
  input_data = sf_model_obs_bcn,
  verbose = FALSE,
  model = "Gau",  # Gaussian model
  miscFitOptions = list(merge.small.bins = FALSE)
)


# Apply Universal Kriging to the whole map
uk_gstat_traffic <- gstat(
  formula = pollutant_obs ~ pollutant_model + buffer_500m,  # Include traffic in the kriging model
  data = sf_model_obs_bcn,  # Data that includes traffic and pollutants
  model = pollutant.fit_loocv_traffic$var_model
)
 
colnames(combined_data)[1] <- "pollutant_model" 
colnames(combined_data)[3] <- "buffer_500m"
# Predict the values of the model in the refined grid
refined_corrected_model_traffic <- predict(object = uk_gstat_traffic, newdata = combined_data)

# Visualize the corrected model using pollutant predictions
mapview::mapview(refined_corrected_model_traffic, zcol = "var1.pred")  # Predictions
mapview::mapview(refined_corrected_model_traffic, zcol = "var1.var")   # Prediction variance

