# Load necessary libraries
library(caret)
library(randomForest)
library(dplyr)
library(ggplot2)
library(naivebayes)  # For Naive Bayes
library(gbm)         # For GBM
library(pROC)        # For ROC and AUC

# Load and preprocess data
data <- read.csv("C:/Users/xuw12/OneDrive - Texas State University/Fall 2024/CE 7393/HW4/CA_AV_CrashReportsSample_v2.csv")
data <- subset(data, select = -c(CrashNarrative))  # Remove CrashNarrative column

# Handle missing values
data <- na.omit(data)

# Separate predictors and target
predictors <- subset(data, select = -Pre_Crash_Mode)  # Exclude Pre_Crash_Mode from predictors
target <- data$Pre_Crash_Mode  # Save target variable

# Encode categorical predictors only
dummies <- dummyVars("~ .", data = predictors)
predictors_encoded <- predict(dummies, newdata = predictors)
predictors_encoded <- as.data.frame(predictors_encoded)

# Combine encoded predictors and target back into one dataset
data_encoded <- cbind(predictors_encoded, Pre_Crash_Mode = target)

# Split data into training and testing sets
set.seed(123)
trainIndex <- createDataPartition(data_encoded$Pre_Crash_Mode, p = 0.7, list = FALSE)
trainData <- data_encoded[trainIndex,]
testData <- data_encoded[-trainIndex,]

# Define control for cross-validation
control <- trainControl(method = "cv", number = 5, classProbs = TRUE, summaryFunction = defaultSummary)

# Initialize results data frame
results <- data.frame(Model = character(), Accuracy = numeric(), Precision = numeric(), Recall = numeric(), F1 = numeric(), AUC = numeric(), stringsAsFactors = FALSE)

# Function to calculate evaluation metrics
evaluate_model <- function(pred, prob, true_labels) {
  conf <- confusionMatrix(pred, true_labels)
  accuracy <- conf$overall['Accuracy']
  precision <- mean(conf$byClass['Pos Pred Value'], na.rm = TRUE)
  recall <- mean(conf$byClass['Sensitivity'], na.rm = TRUE)
  f1 <- 2 * (precision * recall) / (precision + recall)
  auc <- auc(true_labels, prob[,2])  # Calculate AUC for the second class
  return(c(accuracy, precision, recall, f1, auc))
}

# Load necessary libraries
library(caret)
library(text2vec)
library(glmnet)  # For logistic regression
library(lime)    # For explainability
library(wordcloud)
library(tm)
library(ggplot2)

# Load the data
data <- read.csv("C:/Users/xuw12/OneDrive - Texas State University/Fall 2024/CE 7393/HW4/CA_AV_CrashReportsSample_v2.csv")
text_data <- data[c("CrashNarrative", "Pre_Crash_Mode")]

# Ensure no missing values in the relevant columns
text_data <- na.omit(text_data)

# Preprocess text data
text_data$CrashNarrative <- tolower(text_data$CrashNarrative)  # Lowercase the text
text_data$CrashNarrative <- gsub("[^[:alnum:] ]", " ", text_data$CrashNarrative)  # Remove punctuation

# Split the data into training and testing sets
set.seed(123)
trainIndex <- createDataPartition(text_data$Pre_Crash_Mode, p = 0.7, list = FALSE)
trainData <- text_data[trainIndex, ]
testData <- text_data[-trainIndex, ]

# Text Vectorization using TF-IDF
it <- itoken(trainData$CrashNarrative, progressbar = FALSE)
vocab <- create_vocabulary(it)
vectorizer <- vocab_vectorizer(vocab)
dtm_train <- create_dtm(it, vectorizer)

# Create test data matrix
it_test <- itoken(testData$CrashNarrative, progressbar = FALSE)
dtm_test <- create_dtm(it_test, vectorizer)

# Convert dtm_train and dtm_test to data frames for compatibility with lime
dtm_train_df <- as.data.frame(as.matrix(dtm_train))
dtm_test_df <- as.data.frame(as.matrix(dtm_test))

# Train a logistic regression model on the TF-IDF features
model <- cv.glmnet(dtm_train, as.factor(trainData$Pre_Crash_Mode), family = "multinomial")

# ---- Training Process Plots ----

# 1. Cross-Validation MSE Plot
plot(model) + theme_bw(base_size = 30)  # Shows cross-validated MSE across lambda values

# 2. Coefficient Path Plot
plot(model$glmnet.fit, xvar = "lambda", label = TRUE) + theme_bw(base_size = 30)  # Shows the coefficient paths across lambda

# 3. Custom Cross-Validation MSE Plot
# Extract cross-validated mean error and standard deviation
cv_results <- data.frame(
  lambda = model$lambda,
  mean_mse = model$cvm,
  mse_std_error = model$cvsd
)

# Plot cross-validated mean MSE with error bars across lambda values
ggplot(cv_results, aes(x = log(lambda), y = mean_mse)) +
  geom_line(color = "blue") +
  geom_point() +
  geom_errorbar(aes(ymin = mean_mse - mse_std_error, ymax = mean_mse + mse_std_error), width = 0.1) +
  labs(title = "Cross-Validation Mean Squared Error across Lambda",
       x = "Log(Lambda)",
       y = "Mean Squared Error") +
  theme_minimal() + theme_bw(base_size = 22)

# ---- Additional Analysis ----

# Predict on test data
pred <- predict(model, dtm_test, s = "lambda.min", type = "class")
conf_matrix <- confusionMatrix(as.factor(pred), as.factor(testData$Pre_Crash_Mode))
print("Logistic Regression Results for Text Data:")

print(conf_matrix)

# ---- Additional Plots ----

# 1. Word Cloud of Important Words in CrashNarrative
text_corpus <- Corpus(VectorSource(text_data$CrashNarrative))
text_corpus <- tm_map(text_corpus, content_transformer(tolower))
text_corpus <- tm_map(text_corpus, removePunctuation)
text_corpus <- tm_map(text_corpus, removeNumbers)
text_corpus <- tm_map(text_corpus, removeWords, stopwords("en"))

# Generate word cloud
wordcloud(text_corpus, max.words = 100, random.order = FALSE, colors = brewer.pal(8, "Dark2"))

# 2. Confusion Matrix Heatmap
conf_matrix_table <- as.data.frame(conf_matrix$table)
ggplot(conf_matrix_table, aes(x = Prediction, y = Reference, fill = Freq)) +
  geom_tile() +
  scale_fill_gradient(low = "white", high = "blue") +
  labs(title = "Confusion Matrix Heatmap", x = "Predicted", y = "Actual") +
  theme_minimal() + theme_bw(base_size = 22)

# ---- Explainable AI with LIME ----

# Custom model_type function for glmnet
model_type.cv.glmnet <- function(x, ...) {
  return("classification")
}

# Custom predict_model function for glmnet
predict_model.cv.glmnet <- function(x, newdata, type, ...) {
  # Get probabilities for each class
  pred_probs <- predict(x, as.matrix(newdata), s = "lambda.min", type = "response")
  colnames(pred_probs) <- levels(as.factor(trainData$Pre_Crash_Mode))  # Ensure column names match class labels
  as.data.frame(pred_probs)
}

# Register the custom methods with lime
class(model) <- c("cv.glmnet", class(model))  # Explicitly set the model class
explainer <- lime(dtm_train_df, model, bin_continuous = TRUE)

# Explain the predictions for the first 5 instances in the test set
explanation <- explain(dtm_test_df[1:5, ], explainer, n_labels = 1, n_features = 5)

# Plot the explanations
plot_features(explanation) + theme_bw(base_size = 16) +
  labs(title = "Feature Contributions for Predictions")