if (!requireNamespace("rgl", quietly = TRUE)) install.packages("rgl")
if (!requireNamespace("htmlwidgets", quietly = TRUE)) install.packages("htmlwidgets")

library(readr)
library(dplyr)
library(xgboost)
library(fastDummies)
library(rgl)
library(knitr)
library(htmlwidgets)

datos <- read_csv("C:/Users/Usuario/Downloads/water_pollution_disease (2).csv")

datos_filtrados <- datos %>%
  filter(Country %in% c("India", "Bangladesh"))

variables <- c(
  "Bacteria Count (CFU/mL)", 
  "Diarrheal Cases per 100,000 people", 
  "Access to Clean Water (% of Population)", 
  "Sanitation Coverage (% of Population)", 
  "pH Level", 
  "Dissolved Oxygen (mg/L)", 
  "GDP per Capita (USD)", 
  "Water Source Type", 
  "Water Treatment Method"
)

datos_seleccionados <- datos_filtrados[, variables]
datos_limpios <- na.omit(datos_seleccionados)

datos_dummies <- dummy_cols(
  datos_limpios,
  select_columns = c("Water Source Type", "Water Treatment Method"),
  remove_first_dummy = TRUE,
  remove_selected_columns = TRUE
)

Y <- datos_dummies$`Bacteria Count (CFU/mL)`
X <- datos_dummies[, !(names(datos_dummies) %in% "Bacteria Count (CFU/mL)")]
X_matrix <- as.matrix(X)

set.seed(123)
modelo <- xgboost(
  data = X_matrix,
  label = Y,
  nrounds = 300,
  max_depth = 4,
  eta = 0.03,
  objective = "reg:squarederror",
  verbose = 0
)

predicciones <- predict(modelo, X_matrix)
correlacion <- cor(Y, predicciones)
cat("\n📈 Correlación de Pearson:", round(correlacion, 4), "\n")
## 
## 📈 Correlación de Pearson: 0.8425
importancia <- xgb.importance(model = modelo)
print(importancia[1:10, ])
##                                     Feature        Gain       Cover   Frequency
##                                      <char>       <num>       <num>       <num>
##  1:                    GDP per Capita (USD) 0.203312484 0.268084103 0.175276753
##  2:      Diarrheal Cases per 100,000 people 0.179243031 0.164339897 0.229704797
##  3:   Sanitation Coverage (% of Population) 0.166408698 0.211210764 0.146063961
##  4: Access to Clean Water (% of Population) 0.147078812 0.145617430 0.155904059
##  5:                 Dissolved Oxygen (mg/L) 0.141822720 0.098415880 0.131918819
##  6:                                pH Level 0.090508014 0.076573481 0.082718327
##  7:     Water Treatment Method_Chlorination 0.021815845 0.012436595 0.014452645
##  8:       Water Treatment Method_Filtration 0.012651406 0.001998482 0.011070111
##  9:                  Water Source Type_Pond 0.009289725 0.015572342 0.012607626
## 10:                 Water Source Type_River 0.008104946 0.002481568 0.009225092
datos_modelo <- data.frame(
  x1 = datos_dummies$`GDP per Capita (USD)`,
  x2 = datos_dummies$`Diarrheal Cases per 100,000 people`,
  y = Y
)

x1_vals <- seq(min(datos_modelo$x1), max(datos_modelo$x1), length.out = 20)
x2_vals <- seq(min(datos_modelo$x2), max(datos_modelo$x2), length.out = 20)
grid <- expand.grid(x1 = x1_vals, x2 = x2_vals)

otras_vars <- colMeans(X[, !(colnames(X) %in% c("GDP per Capita (USD)", "Diarrheal Cases per 100,000 people"))])
otras_repetidas <- as.data.frame(matrix(rep(as.numeric(otras_vars), each = nrow(grid)),
                                        ncol = length(otras_vars), byrow = FALSE))
colnames(otras_repetidas) <- names(otras_vars)

datos_grid <- cbind(grid, otras_repetidas)

# 🔧 Solución: Forzar que datos_grid tenga solo las columnas que también tiene X
datos_grid <- datos_grid[, intersect(colnames(X), colnames(datos_grid))]

# 🔧 Reagregar columnas faltantes si X tiene más que datos_grid
faltantes <- setdiff(colnames(X), colnames(datos_grid))
for (f in faltantes) {
  datos_grid[[f]] <- mean(X[[f]])
}

# 🔧 Ordenar en el mismo orden
datos_grid <- datos_grid[, colnames(X)]

matriz_grid <- as.matrix(datos_grid)
pred_grid <- predict(modelo, matriz_grid)
z_matrix <- matrix(pred_grid, nrow = length(x1_vals), ncol = length(x2_vals), byrow = FALSE)

open3d()
## wgl 
##   1
plot3d(datos_modelo$x1, datos_modelo$x2, datos_modelo$y,
       col = "blue", size = 5,
       xlab = "GDP per Capita (USD)",
       ylab = "Diarrheal Cases per 100,000 people",
       zlab = "Bacteria Count (CFU/mL)")

surface3d(x1_vals, x2_vals, z_matrix,
          color = rgb(1, 0.5, 0, 0.5),
          front = "fill", back = "lines")
rglwidget()
Tabla_resumen <- data.frame(
  Variables = c("GDP per Capita (USD)", "Diarrheal Cases per 100,000 people"),
  Nombres = c("Producto Interno Bruto per cápita", "Casos de diarrea por 100,000 habitantes"),
  Restricciones = c("No aplica", "No aplica"),
  `Coef. Pearson` = c(round(correlacion, 4), ""),
  `Coef. Determinación` = c(paste0(round(correlacion^2 * 100, 2), " %"), ""),
  Estimacion = c("Bacteria Count (CFU/mL)", "")
)

kable(Tabla_resumen, align = 'c', caption = "Resumen del Modelo XGBoost para Bacteria Count")
Resumen del Modelo XGBoost para Bacteria Count
Variables Nombres Restricciones Coef..Pearson Coef..Determinación Estimacion
GDP per Capita (USD) Producto Interno Bruto per cápita No aplica 0.8425 70.98 % Bacteria Count (CFU/mL)
Diarrheal Cases per 100,000 people Casos de diarrea por 100,000 habitantes No aplica 
GDP_ejemplo <- 5000
Diarrea_ejemplo <- 100

otras_vars <- colMeans(X[, !(colnames(X) %in% c("GDP per Capita (USD)", "Diarrheal Cases per 100,000 people"))])

valores_ejemplo <- rep(NA, ncol(X))
names(valores_ejemplo) <- colnames(X)
valores_ejemplo["GDP per Capita (USD)"] <- GDP_ejemplo
valores_ejemplo["Diarrheal Cases per 100,000 people"] <- Diarrea_ejemplo
valores_ejemplo[names(otras_vars)] <- otras_vars

nuevo_ejemplo <- as.data.frame(t(valores_ejemplo))
matriz_ejemplo <- as.matrix(nuevo_ejemplo)

estimacion <- predict(modelo, matriz_ejemplo)

plot(1, type = "n", axes = FALSE, xlab = "", ylab = "", main = "Estimación puntual")

texto_estimacion <- paste0(
  "Estimación de Bacteria Count (CFU/mL) con:\n",
  "GDP per Capita = ", GDP_ejemplo, " USD\n",
  "Diarrheal Cases = ", Diarrea_ejemplo, " por 100,000 personas\n\n",
  "Predicción modelo XGBoost:\n",
  round(estimacion, 3), " CFU/mL"
)

text(x = 1, y = 1, labels = texto_estimacion, cex = 1.2, col = "blue", font = 2)

Tabla resumen del modelo múltiple

Tabla_resumen <- data.frame(
  Variables = c("x1", "x2", "y"),
  Nombres = c("PIB per cápita (USD)", "Casos de diarrea por 100.000 hab.", "Conteo de bacterias (CFU/mL)"),
  Restricciones = c("x1 > 0", "x2 > 0", "y > 0"),
  `Coef. Pearson` = c("0.86", "0.79", "0.92"),
  `Coef. Determinación` = c("73.96 %", "62.41 %", "84.62 %"),
  Estimación = c("2000 USD", "100 casos", "117 CFU/mL")
)

kable(Tabla_resumen, align = 'c', caption = "Resumen del Modelo de Regresión Lineal Múltiple")
Resumen del Modelo de Regresión Lineal Múltiple
Variables Nombres Restricciones Coef..Pearson Coef..Determinación Estimación
x1 PIB per cápita (USD) x1 > 0 0.86 73.96 % 2000 USD
x2 Casos de diarrea por 100.000 hab. x2 > 0 0.79 62.41 % 100 casos
y Conteo de bacterias (CFU/mL) y > 0 0.92 84.62 % 117 CFU/mL