# ============================================================
# STEP 1: Load Required Packages
# ============================================================

# These packages support modeling, visualization, and data handling
library(dplyr)
library(randomForest)
library(ggplot2)
library(viridis)
library(mgcv)
library(xgboost)
library(Metrics)
library(patchwork)
library(readxl)
library(iml)
library(reshape2)
library(tidyr)
library(ggrepel)
library(openair)
library(MASS)
library(tibble)
library(pdp)
library(plotly)
library(scales)

# ============================================================
# STEP 2: Load Dataset
# ============================================================
as_data <- read_excel("Dataset.xlsx")
View(as_data)
summary(as_data)

# Reshape for plotting
as_data_long <- as_data %>%
  pivot_longer(cols = c(Clay, pH, CEC, OC, TMC),
               names_to = "Parameter",
               values_to = "Value")

#color palette for box fills
param_colors <- c(
  "CEC"  = "#8da0cb",  
  "Clay" = "#66c2a5",  
  "OC"   = "#e78ac3",  
  "pH"   = "#fc8d62",  
  "TMC"  = "#a6d854"   
)

ggplot(as_data_long, aes(x = Parameter, y = Value, fill = Parameter)) +
  geom_boxplot(outlier.shape = NA, color = "black", width = 0.6, alpha = 0.85) +
  geom_jitter(color = "black", width = 0.2, alpha = 0.4, size = 1.8) +
  scale_y_log10(
    breaks = scales::trans_breaks("log10", function(x) 10^x),
    labels = scales::trans_format("log10", math_format(10^.x)),
    limits = c(1, 50000)
  ) +
  scale_fill_manual(values = param_colors) + 
  labs(
    x = "Parameter",
    y = expression("Value (log"[10]*" scale)")
  ) +
  theme_classic(base_size = 14) +
  theme(
    legend.position = "none",
    axis.text = element_text(color = "black"),
    axis.title = element_text(face = "bold"),
    plot.title = element_text(size = 16, face = "bold", hjust = 0.5)
  )


# ============================================================
# STEP 3: Train-Test Split
# ============================================================

set.seed(123)
train_indices <- sample(1:nrow(as_data), size = 0.8 * nrow(as_data))
as_train <- as_data[train_indices, ]
as_test  <- as_data[-train_indices, ]

# ============================================================
# STEP 4: Model CEC as a Function of OC, pH, and Clay
# ============================================================

cec_model <- glm(CEC ~ OC + pH + Clay, data = as_data)
summary(cec_model)


# ============================================================
# STEP 5: Train Prediction Models
# ============================================================

set.seed(42)
# Random Forest
rf_model <- randomForest(PF1 ~ OC + pH + CEC + Clay + TMC, data = as_train, ntree = 500)


# Generalized Additive Model (GAM)
gam_model <- gam(PF1 ~ s(OC) + s(pH) + s(CEC) + s(Clay) + s(TMC), data = as_train)

# XGBoost
x_vars <- c("OC", "pH", "CEC", "Clay", "TMC")
X_train <- as.matrix(as_train[, x_vars])
X_test <- as.matrix(as_test[, x_vars])
y_train <- as_train$PF1
y_test <- as_test$PF1
dtrain <- xgb.DMatrix(data = X_train, label = y_train)
dtest  <- xgb.DMatrix(data = X_test)

xgb_model <- xgboost(
  data = dtrain,
  objective = "reg:squarederror",
  nrounds = 100,
  max_depth = 4,
  eta = 0.1,
  verbosity = 0
)


# ============================================================
# STEP 5A: Variable Importance and Partial Dependence Plots
# ============================================================

# VARIABLE IMPORTANCE
var_imp <- importance(rf_model)
varImpPlot(rf_model,
           main = "Variable Importance – Random Forest",
           col = "blue")


# PARTIAL DEPENDENCE PLOTS
predictor_rf <- Predictor$new(
  rf_model,
  data = as_train[, c("OC", "pH", "CEC", "Clay", "TMC")],
  y = as_train$PF1
)

features_to_plot <- c("OC", "pH", "CEC", "Clay", "TMC")

# Individual PDPs
for (feat in features_to_plot) {
  pdp <- FeatureEffect$new(predictor_rf, feature = feat, method = "pdp")
  print(pdp$plot() + ggtitle(paste("Partial Dependence –", feat)))
}

#Arrange all PDPs in a grid
pdp_plots <- lapply(features_to_plot, function(f) {
  FeatureEffect$new(predictor_rf, feature = f, method = "pdp")$plot() +
    ggtitle(paste("Partial Dependence –", f))
})

wrap_plots(pdp_plots, ncol = 2)

# ============================================================
# STEP 5B: 3D Partial Dependence Plot (OC vs pH)
# ============================================================

pdp_rf_3d <- partial(rf_model, pred.var = c("OC", "pH"), train = as_train, grid.resolution = 20)

z_matrix <- matrix(pdp_rf_3d$yhat, nrow = length(unique(pdp_rf_3d$OC)), ncol = length(unique(pdp_rf_3d$pH)))

fig <- plot_ly(x = unique(pdp_rf_3d$OC), y = unique(pdp_rf_3d$pH), z = z_matrix, type = "surface")

fig <- fig %>% layout(
  title = list(text = "3D Partial Dependence: OC vs pH"),
  scene = list(
    xaxis = list(title = "OC"),
    yaxis = list(title = "pH"),
    zaxis = list(title = "PF1")
  )
)

fig


# ============================================================
# STEP 6: Define USDA Soil Textural Classes (Clay Content)
# ============================================================

usda_clay_classes <- data.frame(
  Texture = c("Sand", "Loamy Sand", "Sandy Loam", "Loam", "Silt Loam", "Silt",
              "Sandy Clay Loam", "Clay Loam", "Silty Clay Loam",
              "Sandy Clay", "Silty Clay", "Clay"),
  Clay = c(5, 10, 10, 25, 15, 10, 25, 30, 35, 35, 45, 50)
)

# ============================================================
# STEP 7: Create Simulation Grid for OC and pH
# ============================================================

sim_base <- expand.grid(
  OC = seq(0.1, 10, by = 0.2),
  pH = seq(4.5, 8.5, by = 0.2)
)
sim_base$TMC <- 20  # Fixed total arsenic concentration


# ============================================================
# STEP 8: Evaluate and Compare Model Performance
# ============================================================
# Gather predictions for each model
test_preds <- as_test %>%
  mutate(
    RF = predict(rf_model, newdata = as_test),
    GAM = predict(gam_model, newdata = as_test),
    XGB = predict(xgb_model, newdata = xgb.DMatrix(data = as.matrix(as_test[, x_vars])))
  ) %>%
  dplyr::select(PF1, RF, GAM, XGB) %>%
  pivot_longer(cols = -PF1, names_to = "Model", values_to = "Prediction") %>%
  mutate(Residual = PF1 - Prediction)

# Compute R², RMSE, MAE for annotation
metrics_df <- data.frame(
  Model = c("RF", "GAM", "XGB"),
  RMSE = c(rmse(as_test$PF1, test_preds$Prediction[test_preds$Model == "RF"]),
           rmse(as_test$PF1, test_preds$Prediction[test_preds$Model == "GAM"]),
           rmse(as_test$PF1, test_preds$Prediction[test_preds$Model == "XGB"])),
  MAE = c(mae(as_test$PF1, test_preds$Prediction[test_preds$Model == "RF"]),
          mae(as_test$PF1, test_preds$Prediction[test_preds$Model == "GAM"]),
          mae(as_test$PF1, test_preds$Prediction[test_preds$Model == "XGB"])),
  R2 = c(cor(as_test$PF1, test_preds$Prediction[test_preds$Model == "RF"])^2,
         cor(as_test$PF1, test_preds$Prediction[test_preds$Model == "GAM"])^2,
         cor(as_test$PF1, test_preds$Prediction[test_preds$Model == "XGB"])^2)
)

# Create annotation text
metrics_df$label <- paste0("R² = ", round(metrics_df$R2, 2),
                           "
RMSE = ", round(metrics_df$RMSE, 2),
                           "
MAE = ", round(metrics_df$MAE, 2))

ggplot(test_preds, aes(x = PF1, y = Prediction, color = Model)) +
  geom_point(alpha = 0.6) +
  geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "black") +
  facet_wrap(~Model) +
  geom_text(data = metrics_df, aes(x = 50, y = 0, label = label), inherit.aes = FALSE, hjust = 0, size = 4) +
  labs(x = "Observed PF1", y = "Predicted PF1") +
  theme_bw(base_size = 14)+
  theme(legend.position = "none")

# Predict on the test set
pred_rf  <- predict(rf_model, newdata = as_test)
pred_gam <- predict(gam_model, newdata = as_test)
pred_xgb <- predict(xgb_model, newdata = xgb.DMatrix(data = as.matrix(as_test[, x_vars])))

# Calculate RMSE, MAE, R2
model_comparison <- data.frame(
  Model = c("Random Forest", "GAM", "XGBoost"),
  RMSE = c(rmse(as_test$PF1, pred_rf), rmse(as_test$PF1, pred_gam), rmse(as_test$PF1, pred_xgb)),
  MAE = c(mae(as_test$PF1, pred_rf), mae(as_test$PF1, pred_gam), mae(as_test$PF1, pred_xgb)),
  R2 = c(cor(as_test$PF1, pred_rf)^2, cor(as_test$PF1, pred_gam)^2, cor(as_test$PF1, pred_xgb)^2)
)

print(model_comparison)

# Taylor Diagram for visual comparison
taylor_data <- data.frame(obs = as_test$PF1, RF = pred_rf, GAM = pred_gam, XGB = pred_xgb)
taylor_long <- melt(taylor_data, id.vars = "obs", variable.name = "Model")

TaylorDiagram(taylor_long, obs = "obs", mod = "value", group = "Model",
              col = c("red", "blue", "green"))

# ============================================================
# STEP 9: Simulate and Visualize Each Model Separately by Texture
# ============================================================

generate_texture_plots <- function(model_name, fill_opt = "C") {
  plots <- list()
  pf1_vals <- c()  # collect PF1 values across textures
  
  for (i in 1:nrow(usda_clay_classes)) {
    texture <- usda_clay_classes$Texture[i]
    clay_val <- usda_clay_classes$Clay[i]
    
    sim_grid <- sim_base
    sim_grid$Clay <- clay_val
    sim_grid$CEC <- predict(cec_model, newdata = sim_grid)
    
    # Predict PF1 using the selected model
    if (model_name == "RF") {
      sim_grid$PF1 <- predict(rf_model, newdata = sim_grid)
    } else if (model_name == "GAM") {
      sim_grid$PF1 <- predict(gam_model, newdata = sim_grid)
    } else if (model_name == "XGB") {
      sim_grid$PF1 <- predict(xgb_model, newdata = xgb.DMatrix(data = as.matrix(sim_grid[, x_vars])))
    }
    
    pf1_vals <- c(pf1_vals, sim_grid$PF1)  # accumulate PF1 range
    plots[[i]] <- list(data = sim_grid, title = texture)
  }
  
  # Determine global PF1 limits
  global_min <- floor(min(pf1_vals, na.rm = TRUE))
  global_max <- ceiling(max(pf1_vals, na.rm = TRUE))
  
  # Regenerate plots using fixed global scale
  final_plots <- lapply(plots, function(entry) {
    ggplot(entry$data, aes(x = OC, y = pH, fill = PF1)) +
      geom_tile() +
      geom_contour(aes(z = PF1), breaks = 10, color = "red") +
      scale_fill_viridis_c(option = fill_opt, limits = c(global_min, global_max)) +
      labs(title = entry$title, fill = "PF1") +
      theme_minimal(base_size = 10) +
      theme(axis.title = element_blank())
  })
  
  return(final_plots)
}


# Generate model plots
rf_plots  <- generate_texture_plots("RF",  fill_opt = "B")
gam_plots <- generate_texture_plots("GAM", fill_opt = "B")
xgb_plots <- generate_texture_plots("XGB", fill_opt = "B")

# Combine with shared legends and common axes
rf_combined <- (wrap_plots(rf_plots, ncol = 4, guides = "collect") &
                  theme(legend.position = "bottom")) +
  plot_annotation(title = "Random Forest – PF1 Across Soil Textures") &
  xlab("OC (%)") & ylab("pH")

gam_combined <- (wrap_plots(gam_plots, ncol = 4, guides = "collect") &
                   theme(legend.position = "bottom")) +
  plot_annotation(title = "GAM – PF1 Across Soil Textures") &
  xlab("OC (%)") & ylab("pH")

xgb_combined <- (wrap_plots(xgb_plots, ncol = 4, guides = "collect") &
                   theme(legend.position = "bottom")) +
  plot_annotation(title = "XGBoost – PF1 Across Soil Textures") &
  xlab("OC (%)") & ylab("pH")

print(rf_combined)
print(gam_combined)
print(xgb_combined)


# ============================================================
# STEP 10: Extract Median and Range of OC–pH in Safe Zones
# ============================================================

safe_summary_list <- list()

for (i in 1:nrow(usda_clay_classes)) {
  texture <- usda_clay_classes$Texture[i]
  clay_val <- usda_clay_classes$Clay[i]
  
  sim_grid <- sim_base
  sim_grid$Clay <- clay_val
  sim_grid$CEC <- predict(cec_model, newdata = sim_grid)
  sim_grid$PF1_rf <- predict(rf_model, newdata = sim_grid)
  
  safe_zone <- sim_grid %>%
    filter(PF1_rf < 10)
  
  if (nrow(safe_zone) > 0) {
    safe_summary_list[[texture]] <- safe_zone %>%
      summarise(
        Texture = texture,
        Clay = clay_val,
        OC_opt = median(OC),
        pH_opt = median(pH),
        OC_min = min(OC),
        OC_max = max(OC),
        pH_min = min(pH),
        pH_max = max(pH),
        .groups = "drop"
      )
  }
}

safe_summary_df <- bind_rows(safe_summary_list)
print(safe_summary_df)


# ============================================================
# STEP 11: 2D Kernel Density Estimate for OC–pH in Safe Zones
# ============================================================

kde_results <- list()

for (i in 1:nrow(usda_clay_classes)) {
  texture <- usda_clay_classes$Texture[i]
  clay_val <- usda_clay_classes$Clay[i]
  
  sim_grid <- sim_base
  sim_grid$Clay <- clay_val
  sim_grid$CEC <- predict(cec_model, newdata = sim_grid)
  sim_grid$PF1_rf <- predict(rf_model, newdata = sim_grid)
  
  safe_zone <- sim_grid %>%
    filter(PF1_rf < 10)
  
  if (nrow(safe_zone) > 10) {
    dens <- kde2d(safe_zone$OC, safe_zone$pH, n = 100)
    max_idx <- which(dens$z == max(dens$z), arr.ind = TRUE)
    kde_peak <- data.frame(
      Texture = texture,
      Clay = clay_val,
      OC_kde = dens$x[max_idx[1]],
      pH_kde = dens$y[max_idx[2]]
    )
    kde_results[[texture]] <- kde_peak
  }
}

kde_summary_df <- bind_rows(kde_results)
print(kde_summary_df)


# ============================================================
# STEP 12: Individual Texture Plots with PF1, Median, KDE Mode 
# ============================================================
# Create an empty list to store plots
rf_overlay_plots <- list()

for (i in 1:nrow(usda_clay_classes)) {
  texture <- usda_clay_classes$Texture[i]
  clay_val <- usda_clay_classes$Clay[i]
  
  sim_grid <- sim_base
  sim_grid$Clay <- clay_val
  sim_grid$CEC <- predict(cec_model, newdata = sim_grid)
  sim_grid$PF1_rf <- predict(rf_model, newdata = sim_grid)
  
  opt_median <- safe_summary_df %>% filter(Texture == texture)
  opt_kde <- kde_summary_df %>% filter(Texture == texture)
  
  median_point <- opt_median %>% transmute(x = OC_opt, y = pH_opt, Type = "Median")
  kde_point <- opt_kde %>% transmute(x = OC_kde, y = pH_kde, Type = "KDE Mode")
  point_overlay <- bind_rows(median_point, kde_point)
  
  p <- ggplot(sim_grid, aes(x = OC, y = pH, fill = PF1_rf)) +
    geom_tile() +
    geom_contour(aes(z = PF1_rf), breaks = 10, color = "red", linewidth = 0.3) +
    geom_point(data = point_overlay,
               aes(x = x, y = y, shape = Type, color = Type),
               inherit.aes = FALSE,
               size = 3) +
    scale_fill_viridis_c(option = "B", limits = c(0, 60), name = "PF1 (%)") +
    scale_shape_manual(values = c("Median" = 16, "KDE Mode" = 17)) +
    scale_color_manual(values = c("Median" = "yellow", "KDE Mode" = "orange")) +
    labs(title = texture, x = "OC (%)", y = "pH") +
    theme_minimal(base_size = 12) +
    theme(
      plot.title = element_text(face = "bold"),
      axis.title = element_blank()
    )
  
  rf_overlay_plots[[i]] <- p
}

rf_overlay_combined <- (wrap_plots(rf_overlay_plots, ncol = 4, guides = "collect") &
                          theme(legend.position = "bottom")) +
  plot_annotation(title = "Random Forest – PF1 with KDE and Median Reference Points") &
  xlab("OC (%)") & ylab("pH")

# Display
print(rf_overlay_combined)

# ============================================================
# STEP 13: Bootstrap KDE Analysis for All Soil Textures
# ============================================================

set.seed(123)

n_boot <- 1000  # Reduce to 100 for quick testing

# Initialize list for bootstrapped KDE results
kde_bootstrap_results <- list()

for (b in 1:n_boot) {
  boot_data <- as_train[sample(1:nrow(as_train), replace = TRUE), ]
  
  boot_rf_model <- randomForest(PF1 ~ OC + pH + CEC + Clay + TMC,
                                data = boot_data, ntree = 500)
  
  for (i in 1:nrow(usda_clay_classes)) {
    texture <- usda_clay_classes$Texture[i]
    clay_val <- usda_clay_classes$Clay[i]
    
    sim_grid <- sim_base
    sim_grid$Clay <- clay_val
    sim_grid$CEC <- predict(cec_model, newdata = sim_grid)
    sim_grid$PF1_rf <- predict(boot_rf_model, newdata = sim_grid)
    
    safe_zone <- sim_grid %>% filter(PF1_rf < 10)
    
    if (nrow(safe_zone) > 10) {
      dens <- kde2d(safe_zone$OC, safe_zone$pH, n = 100)
      max_idx <- which(dens$z == max(dens$z), arr.ind = TRUE)
      kde_point <- data.frame(
        Texture = texture,
        OC_kde = dens$x[max_idx[1]],
        pH_kde = dens$y[max_idx[2]],
        Iteration = b
      )
      kde_bootstrap_results[[length(kde_bootstrap_results) + 1]] <- kde_point
    }
  }
}

# Combine all bootstrap results
kde_summary_all <- bind_rows(kde_bootstrap_results)

# Fill any missing textures using a single-pass KDE from the original RF model
missing_textures <- setdiff(usda_clay_classes$Texture, unique(kde_summary_all$Texture))

for (texture in missing_textures) {
  clay_val <- usda_clay_classes$Clay[usda_clay_classes$Texture == texture]
  
  sim_grid <- sim_base
  sim_grid$Clay <- clay_val
  sim_grid$CEC <- predict(cec_model, newdata = sim_grid)
  sim_grid$PF1_rf <- predict(rf_model, newdata = sim_grid)
  
  safe_zone <- sim_grid %>% filter(PF1_rf < 10)
  
  if (nrow(safe_zone) > 10) {
    dens <- kde2d(safe_zone$OC, safe_zone$pH, n = 100)
    max_idx <- which(dens$z == max(dens$z), arr.ind = TRUE)
    kde_point <- data.frame(
      Texture = texture,
      OC_kde = dens$x[max_idx[1]],
      pH_kde = dens$y[max_idx[2]],
      Iteration = 0  # Indicates fallback
    )
    kde_summary_all <- bind_rows(kde_summary_all, kde_point)
  }
}

# Step: Summarize KDE Bootstrap Results by Texture Class
summary_table <- kde_summary_all %>%
  group_by(Texture) %>%
  summarise(
    OC_mean = round(mean(OC_kde, na.rm = TRUE), 2),
    OC_2.5  = round(quantile(OC_kde, 0.025, na.rm = TRUE), 2),
    OC_97.5 = round(quantile(OC_kde, 0.975, na.rm = TRUE), 2),
    pH_mean = round(mean(pH_kde, na.rm = TRUE), 2),
    pH_2.5  = round(quantile(pH_kde, 0.025, na.rm = TRUE), 2),
    pH_97.5 = round(quantile(pH_kde, 0.975, na.rm = TRUE), 2),
    .groups = "drop"
  )

print(summary_table)



# Define the USDA texture order
usda_order <- c("Clay", "Clay Loam", "Loam", "Loamy Sand", "Sand", "Sandy Clay",
                "Sandy Clay Loam", "Sandy Loam", "Silt", "Silt Loam", "Silty Clay", "Silty Clay Loam")

# Create the summary table
summary_table <- tibble::tibble(
  Texture = c("Clay", "Clay Loam", "Loam", "Loamy Sand", "Sand", "Sandy Clay",
              "Sandy Clay Loam", "Sandy Loam", "Silt", "Silt Loam", "Silty Clay", "Silty Clay Loam"),
  OC_2.5 = c(1.36, 1.39, 1.33, 1.46, 1.54, 1.4, 1.33, 1.46, 1.46, 1.39, 1.36, 1.4),
  OC_97.5 = c(2.87, 3.47, 2.77, 7.63, 7.33, 3.56, 2.77, 7.63, 7.63, 3.15, 2.87, 3.56),
  pH_2.5 = c(5.32, 5.37, 5.39, 5.43, 5.35, 5.41, 5.39, 5.43, 5.43, 5.43, 5.35, 5.41),
  pH_97.5 = c(6.40, 6.24, 6.32, 6.84, 6.73, 6.36, 6.32, 6.84, 6.84, 6.72, 6.32, 6.36)
)

# Calculate ranges and sensitivity
sensitivity_table <- summary_table %>%
  mutate(
    OC_range = OC_97.5 - OC_2.5,
    pH_range = pH_97.5 - pH_2.5,
    Sensitivity_Score = OC_range + pH_range
  ) %>%
  mutate(Texture = factor(Texture, levels = usda_order)) %>%
  arrange(Texture)

# View the reordered sensitivity table
print(sensitivity_table)

# =========================================================================================================
# =========================================================================================================