failure_counts <- labels %>%
  count(failure) %>%
  mutate(pct = round(n / sum(n) * 100, 1))

print(failure_counts)

ggplot(failure_counts, aes(x = reorder(failure, -n), y = n, fill = failure)) +
  geom_col(width = 0.6, show.legend = FALSE) +
  geom_text(aes(label = paste0(n, " (", pct, "%)")), vjust = -0.4, size = 4, color = "black") +
  scale_fill_manual(values = c("Working" = "#2ecc71", "Heating" = "#e67e22", "Bearings" = "#e74c3c")) +
  labs(title = "Session Count by Failure Type", x = "Failure Type", y = "Count") +
  theme_minimal(base_size = 13)

ggplot(labels, aes(x = factor(device_id), fill = failure)) +
  geom_bar(position = "dodge", width = 0.6) +
  scale_fill_manual(values = c("Working" = "#2ecc71", "Heating" = "#e67e22", "Bearings" = "#e74c3c")) +
  labs(title = "Failure Type by Device ID", x = "Device ID", y = "Session Count", fill = "Failure") +
  theme_minimal(base_size = 13)

labels %>%
  count(program, temperature, spin, load) %>%
  arrange(desc(n)) %>%
  head(10) %>%
  knitr::kable(caption = "Top 10 Operating Condition Combinations")

plot_sample %>%
  filter(!is.na(failure)) %>%
  ggplot(aes(x = Current, fill = failure)) +
  geom_density(alpha = 0.5) +
  scale_fill_manual(values = c("Working"="#2ecc71","Heating"="#e67e22","Bearings"="#e74c3c")) +
  labs(title = "Current Signal Density by Failure Type (sampled)",
       x = "Current (raw ADC)", y = "Density", fill = "Failure") +
  theme_minimal(base_size = 13)

plot_sample %>%
  filter(!is.na(failure)) %>%
  ggplot(aes(x = Vibration, fill = failure)) +
  geom_density(alpha = 0.5) +
  scale_fill_manual(values = c("Working"="#2ecc71","Heating"="#e67e22","Bearings"="#e74c3c")) +
  labs(title = "Vibration Signal Density by Failure Type (sampled)",
       x = "Vibration (raw ADC)", y = "Density", fill = "Failure") +
  theme_minimal(base_size = 13)

plot_sample %>%
  filter(!is.na(failure)) %>%
  pivot_longer(cols = c(Current, Vibration), names_to = "sensor", values_to = "value") %>%
  ggplot(aes(x = failure, y = value, fill = failure)) +
  geom_boxplot(outlier.size = 0.3, outlier.alpha = 0.2) +
  scale_fill_manual(values = c("Working"="#2ecc71","Heating"="#e67e22","Bearings"="#e74c3c")) +
  facet_wrap(~sensor, scales = "free_y") +
  labs(title = "Sensor Signal Distribution by Failure Type",
       x = "Failure Type", y = "Signal Value") +
  theme_minimal(base_size = 13) +
  theme(legend.position = "none")

# session_features_raw already contains all 12 signal features + 10 FFT bands
# Just clean up factors and remove rows with no label
session_features <- session_features_raw %>%
  filter(!is.na(failure)) %>%
  mutate(
    failure = factor(failure),
    load    = factor(load),
    program = factor(program)
  )

cat("Feature matrix:", nrow(session_features), "sessions x",
    ncol(session_features) - 7, "candidate features\n")

library(corrplot)

num_feats <- session_features %>%
  select(mean_current, sd_current, range_current, skew_current, kurt_current,
         mean_vibration, sd_vibration, range_vibration, skew_vibration, kurt_vibration,
         temperature, spin)

cor_mat <- cor(num_feats, use = "complete.obs")
corrplot(cor_mat, method = "color", type = "upper", tl.cex = 0.85,
         title = "Feature Correlation Matrix", mar = c(0,0,1.5,0),
         col = colorRampPalette(c("#e74c3c","white","#2980b9"))(200))

numeric_cols <- c("mean_current","sd_current","range_current","skew_current","kurt_current",
                  "mean_vibration","sd_vibration","range_vibration","skew_vibration","kurt_vibration",
                  "temperature","spin")

anova_results <- lapply(numeric_cols, function(col) {
  f    <- as.formula(paste(col, "~ failure"))
  aov_res <- aov(f, data = session_features)
  p    <- summary(aov_res)[[1]][["Pr(>F)"]][1]
  data.frame(feature = col, p_value = round(p, 5))
})

anova_df <- bind_rows(anova_results) %>%
  mutate(significant = ifelse(p_value < 0.05, "YES", "no")) %>%
  arrange(p_value)

knitr::kable(anova_df, caption = "ANOVA p-values — features that differ by failure type")

# Plot the 4 most significant features
top4 <- anova_df %>% filter(significant == "YES") %>% head(4) %>% pull(feature)

session_features %>%
  select(failure, all_of(top4)) %>%
  pivot_longer(-failure, names_to = "feature", values_to = "value") %>%
  ggplot(aes(x = failure, y = value, fill = failure)) +
  geom_boxplot(outlier.size = 0.5, outlier.alpha = 0.3, show.legend = FALSE) +
  scale_fill_manual(values = c("Working" = "#2ecc71", "Heating" = "#e67e22", "Bearings" = "#e74c3c")) +
  facet_wrap(~feature, scales = "free_y", ncol = 2) +
  labs(title = "Top Discriminating Features by Failure Type",
       x = "Failure Type", y = "Value") +
  theme_minimal(base_size = 12)

# Drop: max_current (r~0.99 with mean), max_vibration (r~0.98 with mean)
# Keep: statistically significant + low multicollinearity
selected_cols <- anova_df %>%
  filter(significant == "YES") %>%
  pull(feature)

cat("Selected numeric features:\n", paste(selected_cols, collapse = "\n "), "\n\n")

cat("Categorical features kept: program, load\n")

model_data <- session_features %>%
  select(file, failure, all_of(selected_cols), program, load) %>%
  na.omit()

cat("Final dataset for modelling:", nrow(model_data), "sessions\n")

set.seed(42)
n         <- nrow(model_data)
train_idx <- sample(seq_len(n), size = floor(0.75 * n))
train_df  <- model_data[ train_idx, ] %>% select(-file)
test_df   <- model_data[-train_idx, ] %>% select(-file)

cat("Train:", nrow(train_df), "| Test:", nrow(test_df), "\n")

table(train_df$failure)

library(nnet)

baseline_model <- multinom(failure ~ ., data = train_df, trace = FALSE)

baseline_preds <- predict(baseline_model, newdata = test_df)
baseline_acc   <- mean(baseline_preds == test_df$failure)
cat("Baseline Multinomial LogReg Accuracy:", round(baseline_acc * 100, 1), "%\n")

library(randomForest)
library(caret)

set.seed(42)

# 5-fold CV to tune mtry (number of features tried at each split)
ctrl <- trainControl(method = "cv", number = 5, classProbs = TRUE,
                     summaryFunction = multiClassSummary, savePredictions = TRUE)

rf_grid <- expand.grid(mtry = c(2, 3, 4, 5))

rf_model <- caret::train(
  failure ~ .,
  data      = train_df,
  method    = "rf",
  trControl = ctrl,
  tuneGrid  = rf_grid,
  ntree     = 500,
  importance= TRUE,
  metric    = "Accuracy"
)

print(rf_model)

cat("\nBest mtry:", rf_model$bestTune$mtry, "\n")

rf_preds <- predict(rf_model, newdata = test_df)
rf_probs <- predict(rf_model, newdata = test_df, type = "prob")

cm <- confusionMatrix(rf_preds, test_df$failure)
print(cm)

cm_df <- as.data.frame(cm$table) %>%
  rename(Predicted = Prediction, Actual = Reference)

ggplot(cm_df, aes(x = Actual, y = Predicted, fill = Freq)) +
  geom_tile(colour = "white", linewidth = 0.8) +
  geom_text(aes(label = Freq), size = 6, fontface = "bold", color = "black") +
  scale_fill_gradient(low = "white", high = "#2980b9") +
  labs(title = "Confusion Matrix — Random Forest (Test Set)",
       x = "Actual Class", y = "Predicted Class", fill = "Count") +
  theme_minimal(base_size = 13) +
  theme(legend.position = "right")

class_stats <- as.data.frame(cm$byClass) %>%
  tibble::rownames_to_column("Class") %>%
  mutate(Class = gsub("Class: ", "", Class)) %>%
  select(Class, Sensitivity, Specificity, `Balanced Accuracy`, `F1`) %>%
  mutate(across(where(is.numeric), ~round(., 3)))

knitr::kable(class_stats, caption = "Per-Class Performance Metrics (Random Forest)")

comparison <- data.frame(
  Model    = c("Multinomial Logistic Regression", "Random Forest (CV-tuned)"),
  Accuracy = c(
    round(baseline_acc * 100, 1),
    round(max(rf_model$results$Accuracy) * 100, 1)
  )
)

knitr::kable(comparison, caption = "Model Accuracy Comparison (%)")

ggplot(comparison, aes(x = Model, y = Accuracy, fill = Model)) +
  geom_col(width = 0.5, show.legend = FALSE) +
  geom_text(aes(label = paste0(Accuracy, "%")), vjust = -0.4, size = 5, color = "black") +
  scale_fill_manual(values = c("#95a5a6", "#2980b9")) +
  coord_cartesian(ylim = c(0, 110)) +
  labs(title = "Model Accuracy Comparison", y = "Accuracy (%)", x = "") +
  theme_minimal(base_size = 13)

# FFT columns were computed in the streaming step — extract from session_features
fft_features <- session_features %>%
  select(file, failure, starts_with("fft_"))

cat("FFT features available for", nrow(fft_features), "sessions\n")

head(fft_features[, c("failure","fft_vib_b1","fft_vib_b2","fft_vib_b3")])

fft_features %>%
  select(failure, starts_with("fft_vib")) %>%
  pivot_longer(-failure, names_to = "band", values_to = "energy") %>%
  ggplot(aes(x = band, y = log1p(energy), fill = failure)) +
  geom_boxplot(outlier.size = 0.4, outlier.alpha = 0.3) +
  scale_fill_manual(values = c("Working" = "#2ecc71", "Heating" = "#e67e22", "Bearings" = "#e74c3c")) +
  labs(title = "Vibration FFT Band Energy by Failure Type (log scale)",
       x = "Frequency Band", y = "log(Band Energy)", fill = "Failure") +
  theme_minimal(base_size = 12)

model_data_fft <- model_data %>%
  left_join(fft_features %>% select(-failure), by = "file") %>%
  select(-file) %>%
  na.omit()

set.seed(42)
fft_idx   <- sample(nrow(model_data_fft), floor(0.75 * nrow(model_data_fft)))
train_fft <- model_data_fft[ fft_idx, ]
test_fft  <- model_data_fft[-fft_idx, ]

rf_fft <- caret::train(
  failure ~ ., data = train_fft, method = "rf",
  trControl = ctrl, tuneGrid = rf_grid, ntree = 500, importance = TRUE, metric = "Accuracy"
)

fft_preds <- predict(rf_fft, newdata = test_fft)
fft_acc   <- mean(fft_preds == test_fft$failure)
cat("RF + FFT Accuracy:", round(fft_acc * 100, 1), "%  |  RF time-domain only:",
    round(max(rf_model$results$Accuracy) * 100, 1), "%\n")

slow_features_raw %>%
  filter(!is.na(failure)) %>%
  ggplot(aes(x = failure, y = mean_power, fill = failure)) +
  geom_boxplot(width = 0.5, show.legend = FALSE) +
  scale_fill_manual(values = c("Working"="#2ecc71","Heating"="#e67e22","Bearings"="#e74c3c")) +
  labs(title = "Active Power (W) by Failure Type — Slow Sensor",
       x = "Failure Type", y = "Mean Active Power (W)") +
  theme_minimal(base_size = 13)

comparison_all <- data.frame(
  Model = c(
    "Baseline LogReg — Fast",
    "RF — Fast (time-domain)",
    "RF — Fast + FFT",
    "RF — Fast + SMOTE",
    "RF — Slow (time-domain)"
  ),
  Accuracy = c(
    round(baseline_acc * 100, 1),
    round(max(rf_model$results$Accuracy) * 100, 1),
    round(fft_acc   * 100, 1),
    round(smote_acc * 100, 1),
    round(slow_acc  * 100, 1)
  )
)

knitr::kable(comparison_all, caption = "Full Model Comparison — All Variants")

ggplot(comparison_all, aes(x = reorder(Model, Accuracy), y = Accuracy, fill = Accuracy)) +
  geom_col(width = 0.6, show.legend = FALSE) +
  geom_text(aes(label = paste0(Accuracy, "%")), hjust = -0.15, size = 4, color = "black") +
  coord_flip(ylim = c(0, 115)) +
  scale_fill_gradient(low = "#aed6f1", high = "#1a5276") +
  labs(title = "All Model Variants — Accuracy Comparison", x = "", y = "Accuracy (%)") +
  theme_minimal(base_size = 12)

export_dir <- "model_export"
dir.create(export_dir, showWarnings = FALSE)
cat("Exporting to:", normalizePath(export_dir), "\n")

# Primary model — RF time-domain features
saveRDS(rf_model,   file.path(export_dir, "rf_timedomain.rds"))

# RF + FFT features
saveRDS(rf_fft,     file.path(export_dir, "rf_fft.rds"))

# RF + SMOTE balanced training
saveRDS(rf_smote,   file.path(export_dir, "rf_smote.rds"))

# Slow-sampling RF
saveRDS(rf_slow,    file.path(export_dir, "rf_slow.rds"))

# Baseline logistic regression
saveRDS(baseline_model, file.path(export_dir, "baseline_logreg.rds"))

cat("Models saved:\n")

list.files(export_dir, pattern = "\\.rds$")

# The SMOTE recipe encodes factor levels and feature order — must be bundled with model
saveRDS(upsample_prep, file.path(export_dir, "smote_recipe_prepped.rds"))
cat("Preprocessing recipe saved.\n")

# Bundle the FFT band function and the full feature extraction pipeline
extract_features <- function(sensor_df, labels_df = NULL) {
  library(e1071)
  library(dplyr)

  feats <- sensor_df %>%
    group_by(file, device_id) %>%
    summarise(
      mean_current    = mean(Current,   na.rm = TRUE),
      sd_current      = sd(Current,     na.rm = TRUE),
      range_current   = diff(range(Current,  na.rm = TRUE)),
      skew_current    = skewness(Current,  na.rm = TRUE),
      kurt_current    = kurtosis(Current,  na.rm = TRUE),
      mean_vibration  = mean(Vibration, na.rm = TRUE),
      sd_vibration    = sd(Vibration,   na.rm = TRUE),
      range_vibration = diff(range(Vibration, na.rm = TRUE)),
      skew_vibration  = skewness(Vibration, na.rm = TRUE),
      kurt_vibration  = kurtosis(Vibration, na.rm = TRUE),
      .groups = "drop"
    )

  if (!is.null(labels_df)) {
    feats <- feats %>%
      left_join(labels_df %>%
                  select(device_id, timestamp_begin, timestamp_end, program, temperature, spin, load, failure),
                by = "device_id")
  }
  feats
}

get_fft_bands <- function(signal, n_bands = 5) {
  n <- length(signal)
  if (n < 4) return(rep(NA_real_, n_bands))
  fft_mag <- Mod(fft(signal - mean(signal, na.rm = TRUE)))[2:(n %/% 2 + 1)]
  band_sz <- max(1L, floor(length(fft_mag) / n_bands))
  sapply(1:n_bands, function(b) {
    idx <- ((b - 1L) * band_sz + 1L) : min(b * band_sz, length(fft_mag))
    sum(fft_mag[idx]^2)
  })
}

save(extract_features, get_fft_bands,
     file = file.path(export_dir, "feature_functions.RData"))
cat("Feature engineering functions saved.\n")

pipeline_meta <- list(
  selected_features  = selected_cols,
  categorical_features = c("program", "load"),
  target             = "failure",
  classes            = c("Working", "Heating", "Bearings"),
  best_model         = "rf_fft",
  trained_on         = Sys.Date(),
  n_train_sessions   = nrow(train_df),
  n_fft_bands        = 5L
)

saveRDS(pipeline_meta, file.path(export_dir, "pipeline_metadata.rds"))
cat("Pipeline metadata saved.\n")

# Reload and run a sanity-check prediction on test set
loaded_model <- readRDS(file.path(export_dir, "rf_timedomain.rds"))
loaded_meta  <- readRDS(file.path(export_dir, "pipeline_metadata.rds"))

check_preds <- predict(loaded_model, newdata = test_df)
check_acc   <- mean(check_preds == test_df$failure)

cat("Reloaded model accuracy (sanity check):", round(check_acc * 100, 1), "%\n")

cat("Target classes:", paste(loaded_meta$classes, collapse = " | "), "\n")

cat("Best model flag:", loaded_meta$best_model, "\n")

# ── HOW TO USE IN A NEW R SCRIPT ──────────────────────────────────────────────

# 1. Load model and helpers
model     <- readRDS("model_export/rf_fft.rds")
meta      <- readRDS("model_export/pipeline_metadata.rds")
load("model_export/feature_functions.RData")

# 2. Read a new sensor CSV
new_raw <- read.csv("path/to/new_session_fast.csv")
new_raw$file      <- "new_session_fast.csv"
new_raw$device_id <- 2L   # set from filename

# 3. Extract time-domain features
new_feats <- extract_features(new_raw)

# 4. Add FFT band features
fft_cur <- get_fft_bands(new_raw$Current)
fft_vib <- get_fft_bands(new_raw$Vibration)
for (i in 1:5) {
  new_feats[[paste0("fft_cur_b", i)]] <- fft_cur[i]
  new_feats[[paste0("fft_vib_b", i)]] <- fft_vib[i]
}

# 5. Predict
prediction <- predict(model, newdata = new_feats)
probs      <- predict(model, newdata = new_feats, type = "prob")

cat("Predicted fault class:", as.character(prediction), "\n")
print(probs)
# ──────────────────────────────────────────────────────────────────────────────

cat("\nAll exported files:\n")

file.info(list.files(export_dir, full.names = TRUE))[, c("size")] %>%
  { data.frame(file = list.files(export_dir), size_kb = round(./1024, 1)) } %>%
  knitr::kable(caption = "Exported Model Files")