############################################################
#                                                          #
#    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             #                                           #
#                                                          #
############################################################

#install.packages("gstat")
#install.packages("sp")
#install.packages("automap")
#install.packages("ggmap")
#install.packages("mapview")
#install.packages("raster")
#install.packages("viridis")
#install.packages("viridisLite")
#install.packages("akima")
#install.packages("ggpubr")
#install.packages("jtools")


# Load libraries
library(gstat)
library(sp)
library(sf)
library(automap)
library(ggmap)
library(mapview)
library(raster)
library(viridis)
library(akima)

#library(AgroR)
library(ggplot2)
library(ggpubr)
library(jtools)
#library(rgeos)



# Setting
wdir <- "/Users/ajulian/Documents/Phd/hands-on/handon_final/" # change this to your working directory
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 

mapview(sp_model)  # special points file
mapview(sp_obs_locations)

# 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


# 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                             #
#                                                          #
############################################################

# Define the domain of the region
lon_domain <- c(3.497162, -0.001190186) # lon max, lon min                                                              
lat_domain <- c(43.02355,40.24870) #lat max, lat min

# Create a grid of points to interpolate the model -> not used in the final version
#box -> rgeos::bbox2SP(lat_domain[1],
#                      lat_domain[2],
#                      lon_domain[1],
#                      lon_domain[2],
#                      proj4string = CRS(crs_latlon)) 

# create a grid of points with a finer resolution

# set the coordinates of the domain
x_coords <- c(lon_domain[1],lon_domain[2],lon_domain[2],lon_domain[1]) 
y_coords <- c(lat_domain[1],lat_domain[1],lat_domain[2],lat_domain[2])


# Create a polygon with the coordinates of the domain
domain_polygon <- sp::Polygon(cbind(x_coords, y_coords))
# Create a Polygon object with the domain polygon
domain_polygons <- sp::Polygons(list(domain_polygon), ID = "A")
# Create a SpatialPolygon object with the domain polygons
domain_spatial_polygons <- sp::SpatialPolygons(list(domain_polygons))
# Set the coordinate reference system for the SpatialPolygons object
sp::proj4string(domain_spatial_polygons) <- sp::CRS(crs_latlon)
# Convert the SpatialPolygons object to an sf object
domain_sf <- sf::st_as_sf(domain_spatial_polygons)
# Transform the sf object to the UTM coordinate reference system
domain_sf <- sf::st_transform(domain_sf, crs_utm)

# Set the spatial resolution for the finer grid (in meters, utm units)
coarse_grid_resolution <- 500

# Create a refined grid with the specified resolution
refined_grid <- sf::st_sf(geom = sf::st_make_grid(domain_sf, cellsize = coarse_grid_resolution, what = "centers"), crs = crs_utm)

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

# OK_finergrid is a data frame with the predicted values of the model in the refined grid

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

# Build a dataframe with model&obs information
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)
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")

# Visualize model & obs
head(df_model_obs)
plot(df_model_obs$pollutant_obs,df_model_obs$pollutant_model,xlab="Observations",ylab="Raw model",xlim=c(0,50),ylim=c(0,50))

# 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 

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

# Fit variogram of the residual regression model between obs and raw CALIOPE
pollutant.fit_loocv <- automap::autofitVariogram(formula=pollutant_obs~pollutant_model,input_data=sp_model_obs, verbose=FALSE,model="Gau", miscFitOptions = list(merge.small.bins = FALSE))
plot(pollutant.fit_loocv)

# 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)

# Set the name of the column to be coherent with the previous steps
colnames(OK_finergrid)[1] <- "pollutant_model" 

# Predict the values of the model in the refined grid
refined_corrected_model  <- predict(object=uk_gstat, newdata=OK_finergrid)

# Visualize the corrected model
mapview::mapview(refined_corrected_model,zcol="var1.pred")
mapview::mapview(refined_corrected_model,zcol="var1.var")

############################################################