# ============================================================================
# CART Modeling on HMDA Data (Stock & Watson) - FIXED VERSION
# Classification and Regression Trees for Mortgage Approval Prediction
# ============================================================================

# Install required packages (if not already installed)
if(!require("pacman")) install.packages("pacman")

pacman::p_load(rpart, rpart.plot, caret, dplyr, haven, 
               randomForest, gbm, ggplot2, pROC, tidyr)

# ============================================================================
# STEP 1: Load and Understand the HMDA Data
# ============================================================================

# Read Stata data
hmda_data <- read_dta("data/hmda_sw.dta")

# Initial exploration
cat("\n=== HMDA Dataset Overview ===\n")
print(dim(hmda_data))
cat("\n=== Variable Names ===\n")
print(colnames(hmda_data))
cat("\n=== Data Structure ===\n")
str(hmda_data)
cat("\n=== First few rows ===\n")
head(hmda_data)
cat("\n=== Summary Statistics ===\n")
summary(hmda_data)

# ============================================================================
# STEP 2: Variable Identification and Recoding - FIXED
# ============================================================================

# Stock & Watson HMDA variable mapping:
# s1 = deny (1=denied, 0=approved) - TARGET VARIABLE
# s2 = pirat (payment to income ratio)
# s3 = hirat (housing expense to income ratio)
# s4 = lvrat (loan to value ratio)
# s5 = chist (consumer credit history: 1-6, 1=best)
# s6 = mhist (mortgage credit history: 1-4, 1=best)
# s7 = phist (public record history: 1-4, 1=best)
# s8 = insurance (1=self-employed, 0=not)
# s9 = condominium (1=condo, 0=not)

# Create cleaned dataset with proper variable mapping
hmda_clean <- hmda_data %>%
  mutate(
    # TARGET VARIABLE - FIXED
    deny = as.factor(s1),  # This is the correct target variable
    
    # Financial ratios - FIXED: Added pirat
    pirat = as.numeric(s2),  # Payment to income ratio
    hirat = as.numeric(s3),  # Housing expense to income ratio
    lvrat = as.numeric(s4),  # Loan to value ratio
    
    # Credit history (as factors for tree methods)
    chist = as.factor(s5),   # Consumer credit history
    mhist = as.factor(s6),   # Mortgage credit history  
    phist = as.factor(s7),   # Public record history
    
    # Other characteristics
    insurance = as.factor(s8),
    condominium = as.factor(s9)
  ) %>%
  select(-starts_with("s"))  # Remove original s1-s9 variables

# Add demographic variables if they exist
if("black" %in% names(hmda_data)) {
  hmda_clean$black <- as.factor(hmda_data$black)
}
if("hispanic" %in% names(hmda_data)) {
  hmda_clean$hispanic <- as.factor(hmda_data$hispanic)
}
if("male" %in% names(hmda_data)) {
  hmda_clean$male <- as.factor(hmda_data$male)
}
if("married" %in% names(hmda_data)) {
  hmda_clean$married <- as.factor(hmda_data$married)
}

# Remove rows with missing target variable - FIXED
hmda_clean <- hmda_clean %>%
  filter(!is.na(deny))

# Handle remaining missing values
cat("\n=== Missing Values Before Cleaning ===\n")
print(colSums(is.na(hmda_clean)))

# Remove rows with any missing values (or use imputation if preferred)
hmda_clean <- na.omit(hmda_clean)

cat("\n=== Dataset After Cleaning ===\n")
cat("Observations:", nrow(hmda_clean), "\n")
cat("Variables:", ncol(hmda_clean), "\n")

# Check class balance
cat("\n=== Target Variable Distribution ===\n")
print(table(hmda_clean$deny))
print(prop.table(table(hmda_clean$deny)))

# ============================================================================
# STEP 3: Exploratory Data Analysis
# ============================================================================

# Summary statistics by approval status
cat("\n=== Summary by Denial Status ===\n")
hmda_clean %>%
  group_by(deny) %>%
  summarise(
    n = n(),
    mean_pirat = mean(pirat, na.rm=TRUE),
    mean_hirat = mean(hirat, na.rm=TRUE),
    mean_lvrat = mean(lvrat, na.rm=TRUE),
    .groups = 'drop'
  ) %>%
  print()

# Visualization: Financial ratios by approval status
if(require(ggplot2)) {
  p1 <- ggplot(hmda_clean, aes(x=deny, y=pirat, fill=deny)) +
    geom_boxplot() +
    labs(title="Payment-to-Income Ratio by Loan Decision",
         x="Denied", y="Payment-to-Income Ratio") +
    theme_minimal()
  print(p1)
  
  p2 <- ggplot(hmda_clean, aes(x=deny, y=lvrat, fill=deny)) +
    geom_boxplot() +
    labs(title="Loan-to-Value Ratio by Loan Decision",
         x="Denied", y="Loan-to-Value Ratio") +
    theme_minimal()
  print(p2)
}

# ============================================================================
# STEP 4: Split Data into Training and Testing Sets
# ============================================================================

set.seed(123)  # For reproducibility

# Create 70-30 train-test split with stratification - FIXED
train_index <- createDataPartition(hmda_clean$deny, p = 0.7, list = FALSE)
train_data <- hmda_clean[train_index, ]
test_data <- hmda_clean[-train_index, ]

cat("\n=== Train-Test Split ===\n")
cat("Training observations:", nrow(train_data), "\n")
cat("Testing observations:", nrow(test_data), "\n")
cat("\nTraining set denial rate:", 
    mean(as.numeric(as.character(train_data$deny))), "\n")
cat("Testing set denial rate:", 
    mean(as.numeric(as.character(test_data$deny))), "\n")

# ============================================================================
# STEP 5: Build CART Model (Single Decision Tree)
# ============================================================================

cat("\n\n=== BUILDING SINGLE DECISION TREE ===\n")

# Build initial full tree
cart_model <- rpart(
  deny ~ .,
  data = train_data,
  method = "class",
  control = rpart.control(
    minsplit = 20,      # Minimum observations for split
    minbucket = 7,      # Minimum observations in leaf
    cp = 0.001,         # Complexity parameter
    maxdepth = 10       # Maximum tree depth
  )
)

# Print tree summary
cat("\n=== Tree Summary ===\n")
print(cart_model)

# Variable importance
cat("\n=== Variable Importance ===\n")
print(cart_model$variable.importance)

# Visualize the full tree
cat("\n=== Decision Tree Visualization ===\n")
rpart.plot(
  cart_model,
  main = "CART Model for HMDA Loan Approval",
  type = 4,           # Draw nodes with labels
  extra = 104,        # Show fitted class, probabilities, percentages
  fallen.leaves = TRUE,
  shadow.col = "gray",
  box.palette = "RdYlGn"
)

# ============================================================================
# STEP 6: Evaluate Initial Model
# ============================================================================

# Make predictions on test set
predictions_cart <- predict(cart_model, test_data, type = "class")
predictions_cart_prob <- predict(cart_model, test_data, type = "prob")[,2]

# Confusion Matrix
cat("\n=== Confusion Matrix - Initial CART Model ===\n")
conf_matrix_cart <- confusionMatrix(
  predictions_cart, 
  test_data$deny,
  positive = "1"
)
print(conf_matrix_cart)

# ROC Curve and AUC
roc_cart <- roc(test_data$deny, predictions_cart_prob, quiet=TRUE)
cat("\n=== ROC AUC - Initial CART Model ===\n")
cat("AUC:", auc(roc_cart), "\n")

plot(roc_cart, main="ROC Curve - CART Model", 
     col="blue", lwd=2)

# ============================================================================
# STEP 7: Prune the Tree using Cross-Validation
# ============================================================================

cat("\n\n=== TREE PRUNING ===\n")

# Check complexity parameter table
cat("\n=== CP Table ===\n")
printcp(cart_model)

# Plot cross-validation results
plotcp(cart_model, main="Cross-Validation Results")

# Get optimal CP value
optimal_cp <- cart_model$cptable[which.min(cart_model$cptable[,"xerror"]),"CP"]
cat("\n=== Optimal CP ===\n")
cat("Optimal CP value:", optimal_cp, "\n")

# Prune the tree
cart_pruned <- prune(cart_model, cp = optimal_cp)

cat("\n=== Pruned Tree Summary ===\n")
print(cart_pruned)

# Visualize pruned tree
rpart.plot(
  cart_pruned,
  main = "Pruned CART Model for HMDA Loan Approval",
  type = 4,
  extra = 104,
  fallen.leaves = TRUE,
  shadow.col = "gray",
  box.palette = "RdYlGn"
)

# Evaluate pruned model
predictions_pruned <- predict(cart_pruned, test_data, type = "class")
predictions_pruned_prob <- predict(cart_pruned, test_data, type = "prob")[,2]

cat("\n=== Confusion Matrix - Pruned CART Model ===\n")
conf_matrix_pruned <- confusionMatrix(
  predictions_pruned,
  test_data$deny,
  positive = "1"
)
print(conf_matrix_pruned)

# ROC for pruned model
roc_pruned <- roc(test_data$deny, predictions_pruned_prob, quiet=TRUE)
cat("\n=== ROC AUC - Pruned CART Model ===\n")
cat("AUC:", auc(roc_pruned), "\n")

# ============================================================================
# STEP 8: Ensemble Methods
# ============================================================================

cat("\n\n=== ENSEMBLE METHODS ===\n")

# ----------------------
# 8A: Random Forest
# ----------------------

cat("\n--- Random Forest Model ---\n")

set.seed(123)
rf_model <- randomForest(
  deny ~ .,
  data = train_data,
  ntree = 500,
  mtry = floor(sqrt(ncol(train_data) - 1)),  # sqrt(p) for classification
  importance = TRUE,
  nodesize = 5
)

cat("\n=== Random Forest Summary ===\n")
print(rf_model)

# Variable Importance Plot
cat("\n=== Variable Importance (Random Forest) ===\n")
print(importance(rf_model))
varImpPlot(rf_model, main="Variable Importance - Random Forest")

# Predictions
predictions_rf <- predict(rf_model, test_data, type = "class")
predictions_rf_prob <- predict(rf_model, test_data, type = "prob")[,2]

# Confusion Matrix
cat("\n=== Confusion Matrix - Random Forest ===\n")
conf_matrix_rf <- confusionMatrix(
  predictions_rf,
  test_data$deny,
  positive = "1"
)
print(conf_matrix_rf)

# ROC AUC
roc_rf <- roc(test_data$deny, predictions_rf_prob, quiet=TRUE)
cat("\n=== ROC AUC - Random Forest ===\n")
cat("AUC:", auc(roc_rf), "\n")

# ----------------------
# 8B: Gradient Boosting - FIXED
# ----------------------

cat("\n--- Gradient Boosting Model ---\n")

# Convert factor to numeric for GBM (0 and 1)
train_data_gbm <- train_data %>%
  mutate(deny_numeric = as.numeric(as.character(deny)))

test_data_gbm <- test_data %>%
  mutate(deny_numeric = as.numeric(as.character(deny)))

set.seed(123)
gbm_model <- gbm(
  deny_numeric ~ . - deny,
  data = train_data_gbm,
  distribution = "bernoulli",  # For binary classification
  n.trees = 1000,
  interaction.depth = 4,
  shrinkage = 0.01,
  bag.fraction = 0.7,
  train.fraction = 0.8,
  cv.folds = 5,
  verbose = FALSE
)

cat("\n=== Gradient Boosting Summary ===\n")
print(summary(gbm_model, plotit=FALSE))

# Find optimal number of trees
best_iter <- gbm.perf(gbm_model, method = "cv", plot.it = TRUE)
cat("\nOptimal number of trees:", best_iter, "\n")

# Predictions
predictions_gbm_prob <- predict(
  gbm_model,
  test_data_gbm,
  n.trees = best_iter,
  type = "response"
)
predictions_gbm <- as.factor(ifelse(predictions_gbm_prob > 0.5, 1, 0))

# Confusion Matrix
cat("\n=== Confusion Matrix - Gradient Boosting ===\n")
conf_matrix_gbm <- confusionMatrix(
  predictions_gbm,
  test_data$deny,
  positive = "1"
)
print(conf_matrix_gbm)

# ROC AUC
roc_gbm <- roc(test_data$deny, predictions_gbm_prob, quiet=TRUE)
cat("\n=== ROC AUC - Gradient Boosting ===\n")
cat("AUC:", auc(roc_gbm), "\n")

# ============================================================================
# STEP 9: Compare All Models
# ============================================================================

cat("\n\n=== MODEL COMPARISON ===\n")

# Create comparison data frame
model_comparison <- data.frame(
  Model = c("Single Tree", "Pruned Tree", "Random Forest", "Gradient Boosting"),
  Accuracy = c(
    conf_matrix_cart$overall["Accuracy"],
    conf_matrix_pruned$overall["Accuracy"],
    conf_matrix_rf$overall["Accuracy"],
    conf_matrix_gbm$overall["Accuracy"]
  ),
  Sensitivity = c(
    conf_matrix_cart$byClass["Sensitivity"],
    conf_matrix_pruned$byClass["Sensitivity"],
    conf_matrix_rf$byClass["Sensitivity"],
    conf_matrix_gbm$byClass["Sensitivity"]
  ),
  Specificity = c(
    conf_matrix_cart$byClass["Specificity"],
    conf_matrix_pruned$byClass["Specificity"],
    conf_matrix_rf$byClass["Specificity"],
    conf_matrix_gbm$byClass["Specificity"]
  ),
  AUC = c(
    auc(roc_cart),
    auc(roc_pruned),
    auc(roc_rf),
    auc(roc_gbm)
  )
)

print(model_comparison)

# Plot ROC curves together
plot(roc_cart, col="blue", lwd=2, main="ROC Curves Comparison")
plot(roc_pruned, col="green", add=TRUE, lwd=2)
plot(roc_rf, col="red", add=TRUE, lwd=2)
plot(roc_gbm, col="purple", add=TRUE, lwd=2)
legend("bottomright", 
       legend=c("Single Tree", "Pruned Tree", "Random Forest", "Gradient Boosting"),
       col=c("blue", "green", "red", "purple"),
       lwd=2)

# ============================================================================
# STEP 10: Business Insights and Interpretation - FIXED
# ============================================================================

cat("\n\n=== BUSINESS INSIGHTS ===\n")

# Extract decision rules from pruned tree - FIXED
if(require(rpart.plot)) {
  cat("\n=== Key Decision Rules (Pruned Tree) ===\n")
  rules <- rpart.rules(cart_pruned, cover=TRUE)
  print(rules)
}

# Feature importance across models
cat("\n=== Top 5 Most Important Variables ===\n")

cat("\nRandom Forest:\n")
rf_importance <- importance(rf_model)[order(importance(rf_model)[,3], decreasing=TRUE),]
print(head(rf_importance, 5))

cat("\nGradient Boosting:\n")
gbm_importance <- summary(gbm_model, plotit=FALSE)
print(head(gbm_importance, 5))

# Analyze specific scenarios - FIXED
cat("\n=== Scenario Analysis ===\n")

# Create example scenarios with all required variables
scenario_base <- data.frame(
  pirat = c(0.25, 0.35, 0.45),
  hirat = c(0.20, 0.25, 0.30),
  lvrat = c(0.70, 0.80, 0.95),
  chist = factor(c(1, 3, 5), levels=levels(train_data$chist)),
  mhist = factor(c(1, 2, 4), levels=levels(train_data$mhist)),
  phist = factor(c(1, 2, 3), levels=levels(train_data$phist)),
  insurance = factor(rep(0, 3), levels=levels(train_data$insurance)),
  condominium = factor(rep(0, 3), levels=levels(train_data$condominium))
)

# Add demographic variables if they exist in training data
scenarios <- scenario_base
if("black" %in% names(train_data)) {
  scenarios$black <- factor(rep(0, 3), levels=levels(train_data$black))
}
if("hispanic" %in% names(train_data)) {
  scenarios$hispanic <- factor(rep(0, 3), levels=levels(train_data$hispanic))
}
if("male" %in% names(train_data)) {
  scenarios$male <- factor(rep(1, 3), levels=levels(train_data$male))
}
if("married" %in% names(train_data)) {
  scenarios$married <- factor(rep(1, 3), levels=levels(train_data$married))
}

cat("\nPredicting denial probability for different scenarios:\n")

# Create scenario predictions data frame - FIXED
scenario_predictions <- data.frame(
  Scenario = c("Low Risk", "Medium Risk", "High Risk")
)

# Add predictions only if scenarios has all required variables
tryCatch({
  scenario_predictions$Single_Tree <- predict(cart_pruned, scenarios, type="prob")[,2]
  scenario_predictions$Random_Forest <- predict(rf_model, scenarios, type="prob")[,2]
  
  # For GBM, create numeric version
  scenarios_gbm <- scenarios
  scenarios_gbm$deny_numeric <- 0  # Placeholder, not used for prediction
  scenario_predictions$Gradient_Boosting <- predict(gbm_model, scenarios_gbm, 
                                                     n.trees=best_iter, type="response")
  
  print(scenario_predictions)
}, error = function(e) {
  cat("Error in scenario prediction:", e$message, "\n")
  cat("Skipping scenario analysis due to missing variables\n")
})

# ============================================================================
# STEP 11: Save Results - FIXED
# ============================================================================

cat("\n\n=== SAVING RESULTS ===\n")

# Create directories if they don't exist
if(!dir.exists("models")) dir.create("models")
if(!dir.exists("results")) dir.create("results")

# Save models
saveRDS(cart_pruned, "models/cart_pruned_model.rds")
saveRDS(rf_model, "models/random_forest_model.rds")
saveRDS(gbm_model, "models/gradient_boosting_model.rds")

# Save comparison results
write.csv(model_comparison, "results/model_comparison.csv", row.names=FALSE)

# Save predictions
test_results <- data.frame(
  actual = test_data$deny,
  pred_cart = predictions_pruned,
  pred_rf = predictions_rf,
  pred_gbm = predictions_gbm,
  prob_cart = predictions_pruned_prob,
  prob_rf = predictions_rf_prob,
  prob_gbm = predictions_gbm_prob
)

write.csv(test_results, "results/test_predictions.csv", row.names=FALSE)

cat("\nAll models and results saved successfully!\n")

# ============================================================================
# STEP 12: Final Summary Report
# ============================================================================

cat("\n\n")
cat("============================================================\n")
cat("                 FINAL ANALYSIS REPORT                      \n")
cat("============================================================\n\n")

cat("Dataset: HMDA (Home Mortgage Disclosure Act) - Stock & Watson\n")
cat("Objective: Predict mortgage application denial\n")
cat("Total Observations:", nrow(hmda_clean), "\n")
cat("Training Set:", nrow(train_data), "\n")
cat("Test Set:", nrow(test_data), "\n\n")

cat("BEST PERFORMING MODEL:", model_comparison$Model[which.max(model_comparison$AUC)], "\n")
cat("Highest AUC:", round(max(model_comparison$AUC), 4), "\n")
cat("Highest Accuracy:", round(max(model_comparison$Accuracy), 4), "\n\n")

cat("KEY FINDINGS:\n")
cat("1. Most important predictors for loan denial:\n")
top_vars <- rownames(head(rf_importance, 3))
for(i in 1:length(top_vars)) {
  cat("   -", top_vars[i], "\n")
}
cat("\n")

cat("2. Model Performance:\n")
cat("   - Ensemble methods (RF, GBM) outperform single trees\n")
cat("   - Random Forest provides best balance of accuracy and interpretability\n")
cat("   - Pruning improves generalization of single trees\n\n")

cat("3. Business Recommendations:\n")
cat("   - Use Random Forest model for automated screening\n")
cat("   - Review borderline cases (probability 0.4-0.6) manually\n")
cat("   - Monitor model performance quarterly and retrain as needed\n")
cat("   - Consider fairness and bias testing if demographic variables used\n\n")

cat("============================================================\n")
cat("                    ANALYSIS COMPLETE                       \n")
cat("============================================================\n")