# Titanic ML Shiny App
# A front end for running machine learning classification models on the Titanic dataset

library(shiny)
library(shinydashboard)
library(tidymodels)
library(titanic)
library(tidyverse)
library(janitor)
library(naniar)
library(discrim)
library(yardstick)

# ============================================
# Data Loading and Preprocessing
# ============================================

# Load and clean data
data(titanic_train)
data(titanic_test)

titanic_train <- titanic_train |> clean_names()
titanic_test <- titanic_test |> clean_names()

# Prepare training data
titanic_train2 <- titanic_train |>
  dplyr::select(-passenger_id, -name, -ticket, -cabin) |>
  mutate(
    survived = as_factor(survived),
    pclass = as_factor(pclass),
    sex = as_factor(sex),
    embarked = as_factor(embarked)
  )

# Prepare test data (for predictions)
titanic_test2 <- titanic_test |>
  dplyr::select(-passenger_id, -name, -ticket, -cabin) |>
  mutate(
    pclass = as_factor(pclass),
    sex = as_factor(sex),
    embarked = as_factor(embarked)
  )

# ============================================
# Pre-train models at startup
# ============================================

set.seed(123)

# Split data
titanic_split <- initial_split(titanic_train2, prop = 0.8)
train_data <- training(titanic_split)
test_data <- testing(titanic_split)

# Create recipe for models that need all variables (Random Forest, Naive Bayes)
full_recipe <- recipe(survived ~ ., data = train_data) |>
  step_impute_mean(age) |>
  step_nzv(all_predictors())

# Recipe for C5.0 and Null model (removes categorical vars as in original)
simple_recipe <- recipe(survived ~ ., data = train_data) |>
  step_rm(pclass, sex, embarked) |>
  step_nzv(all_predictors()) |>
  step_impute_mean(age)

# Recipe with normalization for kNN
knn_recipe <- recipe(survived ~ ., data = train_data) |>
  step_impute_mean(age) |>
  step_nzv(all_predictors()) |>
  step_normalize(all_numeric_predictors())

# Recipe with dummy variables for logistic regression
glm_recipe <- recipe(survived ~ ., data = train_data) |>
  step_impute_mean(age) |>
  step_dummy(pclass, sex, embarked) |>
  step_nzv(all_predictors()) |>
  step_normalize(all_numeric_predictors())

# Prep the data
train_full <- prep(full_recipe) |> bake(train_data)
test_full <- prep(full_recipe) |> bake(test_data)

train_simple <- prep(simple_recipe) |> bake(train_data)
test_simple <- prep(simple_recipe) |> bake(test_data)

train_knn <- prep(knn_recipe) |> bake(train_data)
test_knn <- prep(knn_recipe) |> bake(test_data)

train_glm <- prep(glm_recipe) |> bake(train_data)
test_glm <- prep(glm_recipe) |> bake(test_data)

# Train all models
message("Training models...")

# Model 0: Null Model (uses simple recipe without categoricals)
model_null <- null_model() |>
  set_engine("parsnip") |>
  set_mode("classification") |>
  fit(survived ~ ., data = train_simple)

# Model 1: k-Nearest Neighbors (uses normalized data)
model_knn <- nearest_neighbor(neighbors = 7) |>
  set_engine("kknn") |>
  set_mode("classification") |>
  fit(survived ~ ., data = train_knn)

# Model 2: Boosted C5.0 (uses simple recipe without categoricals)
model_c50 <- boost_tree(trees = 20) |>
  set_engine("C5.0") |>
  set_mode("classification") |>
  fit(survived ~ ., data = train_simple)

# Model 3: Random Forest (Ranger) (uses full data with all variables)
model_ranger <- rand_forest(trees = 100) |>
  set_engine("ranger") |>
  set_mode("classification") |>
  fit(survived ~ ., data = train_full)

# Model 4: Regularized Logistic Regression (uses dummy variables)
model_glm <- logistic_reg(penalty = 0.001, mixture = 0.5) |>
  set_engine("glmnet") |>
  set_mode("classification") |>
  fit(survived ~ ., data = train_glm)

# Model 5: Naive Bayes (uses full data with all variables)
model_nb <- naive_Bayes(Laplace = 1) |>
  set_engine("klaR") |>
  set_mode("classification") |>
  fit(survived ~ ., data = train_full)

message("Models trained successfully!")

# Store models and their test data
models_list <- list(
  "Null Model" = list(model = model_null, test_data = test_simple),
  "k-Nearest Neighbors" = list(model = model_knn, test_data = test_knn),
  "Boosted C5.0" = list(model = model_c50, test_data = test_simple),
  "Random Forest" = list(model = model_ranger, test_data = test_full),
  "Logistic Regression" = list(model = model_glm, test_data = test_glm),
  "Naive Bayes" = list(model = model_nb, test_data = test_full)
)

# Function to get model predictions and metrics
get_model_results <- function(model_info, model_name) {
  model <- model_info$model
  test <- model_info$test_data

  preds <- predict(model, test) |>
    bind_cols(test) |>
    bind_cols(predict(model, test, type = "prob"))

  metrics_df <- preds |>
    metrics(truth = survived, estimate = .pred_class)

  conf_mat_obj <- preds |>
    conf_mat(truth = survived, estimate = .pred_class)

  roc_data <- preds |>
    roc_curve(survived, .pred_1)

  roc_auc_val <- preds |>
    roc_auc(survived, .pred_1)

  list(
    predictions = preds,
    metrics = metrics_df,
    conf_mat = conf_mat_obj,
    roc_curve = roc_data,
    roc_auc = roc_auc_val$.estimate,
    accuracy = metrics_df |> filter(.metric == "accuracy") |> pull(.estimate),
    kappa = metrics_df |> filter(.metric == "kap") |> pull(.estimate)
  )
}

# Pre-compute all model results
model_results <- map(models_list, get_model_results, names(models_list))

# ============================================
# UI Definition
# ============================================

ui <- dashboardPage(
  dashboardHeader(title = "Titanic ML Classification"),

  dashboardSidebar(
    sidebarMenu(
      menuItem("Data", tabName = "data", icon = icon("database")),
      menuItem("Preprocessing", tabName = "preprocess", icon = icon("cogs")),
      menuItem("Models", tabName = "models", icon = icon("chart-bar")),
      menuItem("Comparison", tabName = "comparison", icon = icon("balance-scale")),
      menuItem("Predict", tabName = "predict", icon = icon("question-circle"))
    )
  ),

  dashboardBody(
    tabItems(
      # ============================================
      # Data Tab
      # ============================================
      tabItem(
        tabName = "data",
        fluidRow(
          box(
            title = "Training Data Preview (Head)",
            status = "primary",
            solidHeader = TRUE,
            width = 12,
            DT::dataTableOutput("data_head")
          )
        ),
        fluidRow(
          box(
            title = "Training Data Preview (Tail)",
            status = "primary",
            solidHeader = TRUE,
            width = 12,
            DT::dataTableOutput("data_tail")
          )
        ),
        fluidRow(
          box(
            title = "Missing Values Overview",
            status = "info",
            solidHeader = TRUE,
            width = 6,
            plotOutput("vis_miss_plot")
          ),
          box(
            title = "Missing Values by Variable",
            status = "info",
            solidHeader = TRUE,
            width = 6,
            plotOutput("gg_miss_var_plot")
          )
        ),
        fluidRow(
          box(
            title = "Summary Statistics",
            status = "primary",
            solidHeader = TRUE,
            width = 12,
            verbatimTextOutput("summary_stats")
          )
        )
      ),

      # ============================================
      # Preprocessing Tab
      # ============================================
      tabItem(
        tabName = "preprocess",
        fluidRow(
          box(
            title = "Variable Selection",
            status = "primary",
            solidHeader = TRUE,
            width = 6,
            checkboxGroupInput(
              "predictors",
              "Select Predictors to Include:",
              choices = c("pclass", "sex", "age", "sib_sp", "parch", "fare", "embarked"),
              selected = c("age", "sib_sp", "parch", "fare")
            )
          ),
          box(
            title = "Preprocessing Options",
            status = "primary",
            solidHeader = TRUE,
            width = 6,
            checkboxInput("impute_age", "Impute missing age values (mean)", value = TRUE),
            checkboxInput("normalize", "Normalize numeric predictors", value = FALSE),
            checkboxInput("remove_nzv", "Remove near-zero variance predictors", value = TRUE),
            sliderInput(
              "train_prop",
              "Train/Test Split Ratio:",
              min = 0.5,
              max = 0.9,
              value = 0.8,
              step = 0.05
            )
          )
        ),
        fluidRow(
          box(
            title = "Preprocessed Training Data Sample",
            status = "success",
            solidHeader = TRUE,
            width = 12,
            DT::dataTableOutput("preprocessed_data")
          )
        ),
        fluidRow(
          box(
            title = "Recipe Summary",
            status = "info",
            solidHeader = TRUE,
            width = 12,
            verbatimTextOutput("recipe_summary")
          )
        )
      ),

      # ============================================
      # Models Tab
      # ============================================
      tabItem(
        tabName = "models",
        fluidRow(
          box(
            title = "Select Model",
            status = "primary",
            solidHeader = TRUE,
            width = 4,
            selectInput(
              "model_select",
              "Choose a model:",
              choices = names(models_list),
              selected = "Random Forest"
            ),
            actionButton("refresh_model", "Refresh Model Results", class = "btn-primary")
          ),
          box(
            title = "Model Metrics",
            status = "success",
            solidHeader = TRUE,
            width = 8,
            tableOutput("model_metrics")
          )
        ),
        fluidRow(
          box(
            title = "Confusion Matrix",
            status = "info",
            solidHeader = TRUE,
            width = 6,
            plotOutput("conf_mat_plot")
          ),
          box(
            title = "ROC Curve",
            status = "info",
            solidHeader = TRUE,
            width = 6,
            plotOutput("roc_curve_plot")
          )
        ),
        fluidRow(
          box(
            title = "Predicted Probability Distribution",
            status = "primary",
            solidHeader = TRUE,
            width = 12,
            plotOutput("density_plot")
          )
        )
      ),

      # ============================================
      # Comparison Tab
      # ============================================
      tabItem(
        tabName = "comparison",
        fluidRow(
          box(
            title = "Combined ROC Curves",
            status = "primary",
            solidHeader = TRUE,
            width = 12,
            plotOutput("combined_roc_plot", height = "500px")
          )
        ),
        fluidRow(
          box(
            title = "Model Comparison Table",
            status = "success",
            solidHeader = TRUE,
            width = 12,
            DT::dataTableOutput("comparison_table")
          )
        ),
        fluidRow(
          box(
            title = "Density Plots Comparison",
            status = "info",
            solidHeader = TRUE,
            width = 12,
            plotOutput("combined_density_plot", height = "600px")
          )
        )
      ),

      # ============================================
      # Predict Tab
      # ============================================
      tabItem(
        tabName = "predict",
        fluidRow(
          box(
            title = "Enter Passenger Information",
            status = "primary",
            solidHeader = TRUE,
            width = 6,
            selectInput("pred_pclass", "Passenger Class:", choices = c("1" = 1, "2" = 2, "3" = 3), selected = "1"),
            selectInput("pred_sex", "Sex:", choices = c("male", "female"), selected = "male"),
            numericInput("pred_age", "Age:", value = 30, min = 0, max = 100, step = 1),
            numericInput("pred_sibsp", "Siblings/Spouses Aboard:", value = 0, min = 0, max = 10, step = 1),
            numericInput("pred_parch", "Parents/Children Aboard:", value = 0, min = 0, max = 10, step = 1),
            numericInput("pred_fare", "Fare:", value = 50, min = 0, max = 600, step = 1),
            selectInput("pred_embarked", "Port of Embarkation:", choices = c("C" = "C", "Q" = "Q", "S" = "S"), selected = "S")
          ),
          box(
            title = "Select Model for Prediction",
            status = "primary",
            solidHeader = TRUE,
            width = 6,
            selectInput(
              "predict_model",
              "Choose a model:",
              choices = names(models_list),
              selected = "Random Forest"
            ),
            actionButton("make_prediction", "Predict Survival", class = "btn-primary btn-lg"),
            hr(),
            uiOutput("prediction_result")
          )
        )
      )
    )
  )
)

# ============================================
# Server Logic
# ============================================

server <- function(input, output, session) {

  # ----------------------------------------
  # Data Tab
  # ----------------------------------------

  output$data_head <- DT::renderDataTable({
    DT::datatable(
      head(titanic_train2, 20),
      options = list(scrollX = TRUE, pageLength = 10),
      rownames = FALSE
    )
  })

  output$data_tail <- DT::renderDataTable({
    DT::datatable(
      tail(titanic_train2, 20),
      options = list(scrollX = TRUE, pageLength = 10),
      rownames = FALSE
    )
  })

  output$vis_miss_plot <- renderPlot({
    vis_miss(titanic_train2)
  })

  output$gg_miss_var_plot <- renderPlot({
    gg_miss_var(titanic_train2, show_pct = TRUE)
  })

  output$summary_stats <- renderPrint({
    summary(titanic_train2)
  })

  # ----------------------------------------
  # Preprocessing Tab
  # ----------------------------------------

  preprocessed_recipe <- reactive({
    selected_vars <- input$predictors

    # Build formula based on selected predictors
    if (length(selected_vars) == 0) {
      formula_str <- "survived ~ ."
    } else {
      # We'll include all and use step_rm for excluded ones
      formula_str <- "survived ~ ."
    }

    rec <- recipe(survived ~ ., data = train_data)

    # Remove excluded predictors
    all_vars <- c("pclass", "sex", "age", "sib_sp", "parch", "fare", "embarked")
    excluded_vars <- setdiff(all_vars, selected_vars)
    if (length(excluded_vars) > 0) {
      for (v in excluded_vars) {
        rec <- rec |> step_rm(!!sym(v))
      }
    }

    if (input$impute_age && "age" %in% selected_vars) {
      rec <- rec |> step_impute_mean(age)
    }

    if (input$remove_nzv) {
      rec <- rec |> step_nzv(all_predictors())
    }

    if (input$normalize) {
      rec <- rec |> step_normalize(all_numeric_predictors())
    }

    rec |> prep()
  })

  output$preprocessed_data <- DT::renderDataTable({
    baked_data <- bake(preprocessed_recipe(), train_data)
    DT::datatable(
      head(baked_data, 50),
      options = list(scrollX = TRUE, pageLength = 10),
      rownames = FALSE
    )
  })

  output$recipe_summary <- renderPrint({
    summary(preprocessed_recipe())
  })

  # ----------------------------------------
  # Models Tab
  # ----------------------------------------

  selected_model_results <- reactive({
    model_name <- input$model_select
    model_results[[model_name]]
  })

  output$model_metrics <- renderTable({
    results <- selected_model_results()
    data.frame(
      Metric = c("Accuracy", "Kappa", "ROC AUC"),
      Value = round(c(results$accuracy, results$kappa, results$roc_auc), 4)
    )
  }, striped = TRUE, hover = TRUE, bordered = TRUE)

  output$conf_mat_plot <- renderPlot({
    results <- selected_model_results()
    autoplot(results$conf_mat, type = "mosaic") +
      labs(title = paste("Confusion Matrix -", input$model_select)) +
      theme_minimal()
  })

  output$roc_curve_plot <- renderPlot({
    results <- selected_model_results()
    autoplot(results$roc_curve) +
      labs(title = paste("ROC Curve -", input$model_select)) +
      theme_minimal()
  })

  output$density_plot <- renderPlot({
    results <- selected_model_results()
    results$predictions |>
      ggplot() +
      geom_density(aes(x = .pred_1, fill = survived), alpha = 0.5) +
      labs(
        title = paste("Predicted Probability Distribution -", input$model_select),
        x = "Predicted Probability of Survival",
        y = "Density",
        fill = "Actual Survival"
      ) +
      theme_minimal() +
      scale_fill_manual(values = c("#E41A1C", "#377EB8"))
  })

  # ----------------------------------------
  # Comparison Tab
  # ----------------------------------------

  output$combined_roc_plot <- renderPlot({
    roc_data <- map2_dfr(
      model_results,
      names(model_results),
      ~ .x$roc_curve |> mutate(model = .y)
    )

    ggplot(roc_data, aes(x = 1 - specificity, y = sensitivity, color = model)) +
      geom_path(linewidth = 1) +
      geom_abline(lty = 2, color = "gray50") +
      coord_equal() +
      labs(
        title = "ROC Curves Comparison",
        x = "1 - Specificity (False Positive Rate)",
        y = "Sensitivity (True Positive Rate)",
        color = "Model"
      ) +
      theme_minimal() +
      theme(legend.position = "bottom") +
      scale_color_brewer(palette = "Set1")
  })

  output$comparison_table <- DT::renderDataTable({
    comparison_df <- map2_dfr(
      model_results,
      names(model_results),
      ~ data.frame(
        Model = .y,
        Accuracy = round(.x$accuracy, 4),
        Kappa = round(.x$kappa, 4),
        `ROC AUC` = round(.x$roc_auc, 4),
        check.names = FALSE
      )
    )

    DT::datatable(
      comparison_df,
      options = list(pageLength = 10, order = list(list(2, "desc"))),
      rownames = FALSE
    ) |>
      DT::formatStyle(
        "Accuracy",
        background = DT::styleColorBar(comparison_df$Accuracy, "lightgreen"),
        backgroundSize = "98% 88%",
        backgroundRepeat = "no-repeat",
        backgroundPosition = "center"
      )
  })

  output$combined_density_plot <- renderPlot({
    density_data <- map2_dfr(
      model_results,
      names(model_results),
      ~ .x$predictions |>
        select(.pred_1, survived) |>
        mutate(model = .y)
    )

    ggplot(density_data, aes(x = .pred_1, fill = survived)) +
      geom_density(alpha = 0.5) +
      facet_wrap(~ model, scales = "free_y") +
      labs(
        title = "Predicted Probability Distributions by Model",
        x = "Predicted Probability of Survival",
        y = "Density",
        fill = "Actual Survival"
      ) +
      theme_minimal() +
      scale_fill_manual(values = c("#E41A1C", "#377EB8"))
  })

  # ----------------------------------------
  # Predict Tab
  # ----------------------------------------

  observeEvent(input$make_prediction, {
    model_name <- input$predict_model
    model_info <- models_list[[model_name]]
    model <- model_info$model

    # Create new passenger data
    new_passenger <- data.frame(
      pclass = factor(input$pred_pclass, levels = c(1, 2, 3)),
      sex = factor(input$pred_sex, levels = c("male", "female")),
      age = input$pred_age,
      sib_sp = as.double(input$pred_sibsp),
      parch = as.double(input$pred_parch),
      fare = input$pred_fare,
      embarked = factor(input$pred_embarked, levels = c("C", "Q", "S"))
    )

    # Apply the appropriate recipe based on model
    if (model_name == "k-Nearest Neighbors") {
      new_data <- prep(knn_recipe) |> bake(new_passenger)
    } else if (model_name == "Logistic Regression") {
      new_data <- prep(glm_recipe) |> bake(new_passenger)
    } else if (model_name %in% c("Null Model", "Boosted C5.0")) {
      new_data <- prep(simple_recipe) |> bake(new_passenger)
    } else {
      new_data <- prep(full_recipe) |> bake(new_passenger)
    }

    # Make prediction
    pred_class <- predict(model, new_data)
    pred_prob <- predict(model, new_data, type = "prob")

    # Store results for display
    rv$prediction <- list(
      class = pred_class$.pred_class,
      prob_survived = pred_prob$.pred_1,
      prob_died = pred_prob$.pred_0
    )
  })

  rv <- reactiveValues(prediction = NULL)

  output$prediction_result <- renderUI({
    if (is.null(rv$prediction)) {
      return(div(
        class = "alert alert-info",
        "Click 'Predict Survival' to see the prediction."
      ))
    }

    pred <- rv$prediction
    survival_color <- if (pred$class == 1) "green" else "red"
    survival_text <- if (pred$class == 1) "SURVIVED" else "DID NOT SURVIVE"

    div(
      h3("Prediction Result"),
      div(
        style = paste("background-color:", survival_color, "; color: white; padding: 20px; border-radius: 10px; text-align: center;"),
        h2(strong(survival_text))
      ),
      br(),
      h4("Survival Probability:"),
      div(
        class = "progress",
        div(
          class = "progress-bar",
          style = paste0("width: ", round(pred$prob_survived * 100, 1), "%; background-color: #377EB8;"),
          strong(paste0(round(pred$prob_survived * 100, 1), "%"))
        )
      ),
      br(),
      tableOutput("prediction_table")
    )
  })

  output$prediction_table <- renderTable({
    if (is.null(rv$prediction)) return(NULL)
    pred <- rv$prediction
    data.frame(
      Outcome = c("Did Not Survive", "Survived"),
      Probability = c(
        paste0(round(pred$prob_died * 100, 2), "%"),
        paste0(round(pred$prob_survived * 100, 2), "%")
      )
    )
  })
}

# Run the application
shinyApp(ui = ui, server = server)