library(CAST)
library(sf)
library(terra)
library(dplyr)
library(ggplot2)
library(caret)
library(blockCV)

set.seed(44)

if (!dir.exists("data")) {
  data_url <- "https://jakubnowosad.com/ml4eo2026workshop/data.zip"
  download.file(data_url, "data.zip")
  unzip("data.zip", exdir = ".")
}

# Read the Germany predictor stack and sample points
predictors <- rast(c("data/germany/predictors/bio1.tif", "data/germany/predictors/elevation.tif", "data/germany/predictors/slope.tif"))
samples <- read_sf("data/germany/samples/clustered_samples.gpkg")
response <- "outcome"
predictor_names <- names(predictors)

plot(predictors)
plot(samples)

# Extract predictor values at sample locations and prepare model training data
samples_df <- terra::extract(predictors, samples, ID = FALSE)
samples_df <- cbind(samples, samples_df)
train_data <- samples_df |> 
  st_drop_geometry() |>
  dplyr::select(all_of(predictor_names))

# 1) AoA based only on the training data ------------------------------------------------
# See https://doi.org/10.1111/2041-210X.13650 for details on the method and interpretation of results
aoa_data <- CAST::aoa(
  newdata = predictors,
  train = train_data,
  variables = predictor_names,
  verbose = FALSE
)

# DI -- dissimilarity index -- values close to 0 indicate that the predictor values at a location are similar to those in the training data, while higher values indicate increasing dissimilarity
plot(aoa_data) +
  ggtitle("Distributions of training data and predictors in predictor space (DI values)")

# DI map
plot(aoa_data$DI, main = "Distribution of DI values across the study area")

# AOA map -- values of 1 indicate that the predictor values at a location are within the range of the training data, while values of 0 indicate that they are outside this range
# The threshold for determining whether a location is within the area of applicability is set at the outlier-removed maximum DI value observed in the training data
plot(aoa_data$AOA, main = "Area of Applicability (AOA) across the study area")

# LPD -- local point density
# See https://doi.org/10.5194/gmd-18-10185-2025 for details on the method and interpretation of results
lpd_data <- CAST::aoa(
  newdata = predictors,
  train = train_data,
  variables = predictor_names,
  LPD = TRUE,
  verbose = FALSE
)

# LPD values indicate the density of training data points in the predictor space around a given location, with higher values indicating a higher density of training data points and thus potentially more reliable predictions
plot(lpd_data$LPD, main = "Local Point Density (LPD) across the study area")

# 2) Different validation workflows ------------------------------------------------------
# base model
grid <- expand.grid(
  mtry = 2,
  splitrule = "variance",
  min.node.size = 5
)
rf_base <- train(
  x = train_data,
  y = samples_df[[response]],
  method = "ranger",
  trControl = trainControl(method = "none"),
  tuneGrid = grid,
  metric = "RMSE",
  num.trees = 300,
  importance = "impurity"
)
rf_base_pred <- predict(predictors, model = rf_base, na.rm = TRUE)
plot(rf_base_pred)

# 2a) Random CV
random_folds <- createFolds(samples_df[[response]], k = 5, returnTrain = TRUE)

indexOut <- lapply(random_folds, function(train_idx) setdiff(seq_len(nrow(samples_df)), train_idx))

ctrl_random <- trainControl(
  method = "cv",
  index = random_folds,
  indexOut = indexOut,
  savePredictions = "final"
)

rf_random <- train(
  x = train_data,
  y = samples_df[[response]],
  method = "ranger",
  trControl = ctrl_random,
  tuneGrid = grid,
  metric = "RMSE",
  num.trees = 300,
  importance = "impurity"
)
random_rmse <- sqrt(mean((rf_random$pred$pred - rf_random$pred$obs) ^ 2))
random_rmse

# 2b) Spatial CV
spatial_folds <- blockCV::cv_spatial(
  x = samples_df,
  k = 5,
  seed = 6624
)
blockCV::cv_plot(spatial_folds, samples_df)
train_ids <- lapply(spatial_folds$folds_list, function(x) x[[1]])
test_ids <- lapply(spatial_folds$folds_list, function(x) x[[2]])

ctrl_spatial <- trainControl(
  method = "cv",
  index = train_ids,
  indexOut = test_ids,
  savePredictions = "final",
  verboseIter = FALSE
)

rf_spatial <- train(
  x = train_data,
  y = samples_df[[response]],
  method = "ranger",
  trControl = ctrl_spatial,
  tuneGrid = grid,
  metric = "RMSE",
  num.trees = 300,
  importance = "impurity"
)
spatial_rmse <- sqrt(mean((rf_spatial$pred$pred - rf_spatial$pred$obs) ^ 2))
spatial_rmse

# 2c) kNNDM CV
# https://doi.org/10.5194/gmd-17-5897-2024 for details on the method and interpretation of results
kn <- CAST::knndm(
  tpoints = samples_df,
  modeldomain = predictors[[1]],
  k = 5
)

ctrl_knndm <- trainControl(
  method = "cv",
  index = kn$indx_train,
  indexOut = kn$indx_test,
  savePredictions = "final",
  verboseIter = FALSE
)

rf_knndm <- train(
  x = train_data,
  y = samples_df[[response]],
  method = "ranger",
  trControl = ctrl_knndm,
  tuneGrid = grid,
  metric = "RMSE",
  num.trees = 300,
  importance = "impurity"
)
knndm_rmse <- sqrt(mean((rf_knndm$pred$pred - rf_knndm$pred$obs) ^ 2))
knndm_rmse

# True RMSE 
true_outcome_raster <- rast("data/germany/true_outcome/true_outcome.tif")
plot(true_outcome_raster)

rf_base_pred_values <- terra::values(rf_base_pred, mat = FALSE)
true_outcome_raster_values <- terra::values(true_outcome_raster, mat = FALSE)
true_rmse <- sqrt(mean((rf_base_pred_values - true_outcome_raster_values) ^ 2, na.rm = TRUE))
true_rmse

# Compare
rmse_results <- data.frame(
  Method = c("Random CV", "Spatial CV", "kNNDM CV", "True RMSE"),
  RMSE = c(random_rmse, spatial_rmse, knndm_rmse, true_rmse)
)
rmse_results

# Bonus 2d) kNNDM CV based on the feature space
kn_feature <- CAST::knndm(
  tpoints = samples_df[, predictor_names],
  modeldomain = predictors,
  k = 5,
  dist_space = "feature"
)
ctrl_knndm_feature <- trainControl(
  method = "cv",
  index = kn_feature$indx_train,
  indexOut = kn_feature$indx_test,
  savePredictions = "final",
  verboseIter = FALSE
)
rf_knndm_feature <- train(
  x = train_data,
  y = samples_df[[response]],
  method = "ranger",
  trControl = ctrl_knndm_feature,
  tuneGrid = grid,
  metric = "RMSE",
  num.trees = 300,
  importance = "impurity"
)
knndm_feature_rmse <- sqrt(mean((rf_knndm_feature$pred$pred - rf_knndm_feature$pred$obs) ^ 2))
knndm_feature_rmse

# Compare
rmse_results2 <- data.frame(
  Method = c("Random CV", "Spatial CV", "kNNDM CV", "kNNDM CV (feature space)", "True RMSE"),
  RMSE = c(random_rmse, spatial_rmse, knndm_rmse, knndm_feature_rmse, true_rmse)
)
rmse_results2

# Bonus 2e) geodist
geodist_samples <- CAST::geodist(samples, predictors)
plot(geodist_samples)

geodist_samples_random <- CAST::geodist(samples, predictors, CVtest = ctrl_random$indexOut)
plot(geodist_samples_random)

geodist_samples_spatial <- CAST::geodist(samples, predictors, CVtest = ctrl_spatial$indexOut)
plot(geodist_samples_spatial)

geodist_samples_knndm <- CAST::geodist(samples, predictors, CVtest = ctrl_knndm$indexOut)
plot(geodist_samples_knndm)

# 3) AoA based on the trained models -------------------------------------------------------
aoa_model_random <- aoa(
  newdata = predictors,
  model = rf_random,
  LPD = FALSE,
  verbose = FALSE
)

plot(aoa_model_random) +
  ggtitle("Distributions of training data and predictors in predictor space (DI values)")
plot(aoa_model_random$AOA, main = "Area of Applicability (AOA) across the study area")

aoa_model_knndm <- aoa(
  newdata = predictors,
  model = rf_knndm,
  LPD = FALSE,
  verbose = FALSE
)

plot(aoa_model_knndm) +
  ggtitle("Distributions of training data and predictors in predictor space (DI values)")
plot(aoa_model_knndm$AOA, main = "Area of Applicability (AOA) across the study area")

# 4) Local point density based on the trained models -------------------------------------------------------
lpd_model_random <- CAST::aoa(
  newdata = predictors,
  model = rf_random,
  LPD = TRUE,
  verbose = FALSE
)
plot(lpd_model_random$LPD, main = "Local Point Density (LPD) across the study area")

lpd_model_knndm <- CAST::aoa(
  newdata = predictors,
  model = rf_knndm,
  LPD = TRUE,
  verbose = FALSE
)
plot(lpd_model_knndm$LPD, main = "Local Point Density (LPD) across the study area")

# 5) Error profiles (kNNDM CV) -------------------------------------------------------
errormodel_LPD <- errorProfiles(rf_knndm, lpd_model_knndm, variable = "LPD")
plot(errormodel_LPD) +
  ggtitle("Error profiles based on LPD values")

expected_error_LPD <- terra::predict(lpd_model_knndm$LPD, errormodel_LPD)
plot(expected_error_LPD)

updated_AOA <- lpd_model_knndm$LPD > 50
plot(updated_AOA)

# 6) Relationship between AoA and true error -------------------------------------------------------
true_outcome_raster <- rast("data/germany/true_outcome/true_outcome.tif")
plot(true_outcome_raster)

true_outcome_raster_aoa <- terra::mask(true_outcome_raster, aoa_model_knndm$AOA, maskvalue = 0)
plot(true_outcome_raster_aoa, main = "True outcome masked by AOA")

rf_base_pred_raster_aoa <- terra::mask(rf_base_pred, aoa_model_knndm$AOA, maskvalue = 0)
plot(rf_base_pred_raster_aoa, main = "Predictions masked by AOA")

rf_base_pred_values_aoa <- terra::values(rf_base_pred_raster_aoa, mat = FALSE)
true_outcome_raster_values <- terra::values(true_outcome_raster_aoa, mat = FALSE)
aoa_masked_rmse <- sqrt(mean((rf_base_pred_values_aoa - true_outcome_raster_values) ^ 2, na.rm = TRUE))
aoa_masked_rmse

# Compare
rmse_results3 <- data.frame(
  Method = c("Random CV", "Spatial CV", "kNNDM CV", "kNNDM CV (feature space)", "True RMSE", "AOA-masked RMSE"),
  RMSE = c(random_rmse, spatial_rmse, knndm_rmse, knndm_feature_rmse, true_rmse, aoa_masked_rmse)
)
rmse_results3