library(ggplot2)
library(dplyr)
library(readr)
library(viridisLite)

# Usage: Rscript combined_instr_fault_correlation_heatmap.r exp_abspath1 ... [faults_file]

args <- commandArgs(trailingOnly = TRUE)
if (length(args) < 1) {
  stop("Need at least 1 experiment")
}

csv_suffix <- if (grepl("\\.csv$", args[length(args)])) {
  args[length(args)]
} else {
  "faults.csv"
}
exp_args <- if (grepl("\\.csv$", args[length(args)])) {
  args[-length(args)]
} else {
  args
}

all_data <- data.frame()
all_mnem <- data.frame()

for (arg in exp_args) {
  csv_file <- file.path(arg, csv_suffix)
  if (!file.exists(csv_file)) {
    warning(paste("Missing:", csv_file))
    next
  }

  df <- read_csv(csv_file, col_types = cols())
  df$experiment <- basename(arg)
  all_data <- bind_rows(all_data, df)

  # TODO: This is ignoring any filters currently
  mnem_file <- file.path(arg, "mnemonics.csv")
  if (!file.exists(mnem_file)) {
    warning(paste("Missing:", mnem_file))
    next
  }
  mdf <- read_csv(mnem_file, col_types = cols())
  all_mnem <- bind_rows(all_mnem, mdf)
}

if (nrow(all_data) == 0) {
  stop("No faults.csv data loaded")
}

if (nrow(all_mnem) == 0) {
  stop("No mnemonics.csv data loaded")
}

# no OK_MARKER, sum GROUP1 + TRAP.
all_data <- all_data |>
  filter(!is.na(mnemonic), mnemonic != "NULL") |>
  filter(resulttype != "OK_MARKER") |>
  mutate(
    resulttype = ifelse(resulttype == "GROUP1_MARKER", "TRAP", resulttype)
  )

if (nrow(all_data) == 0) {
  stop("No failure-marker data to plot")
}

# Sum faults per (mnemonic, marker) pair for all experiments
heat <- all_data |>
  group_by(mnemonic, resulttype) |>
  summarise(faults = sum(faults, na.rm = TRUE), .groups = "drop")

# Sum mnemonic counts for all experiments
mnem_counts <- all_mnem |>
  filter(!is.na(mnemonic), mnemonic != "NULL") |>
  group_by(mnemonic) |>
  summarise(count = sum(count, na.rm = TRUE), .groups = "drop")

# Normalize by mnemonic count
heat <- heat |>
  left_join(mnem_counts, by = "mnemonic") |>
  filter(!is.na(count), count > 0) |>
  mutate(fault_rate = faults / count)

if (nrow(heat) == 0) {
  stop("Heat join failed")
}

# Order by fault rate
mnem_order <- heat |>
  group_by(mnemonic) |>
  summarise(total = sum(fault_rate), .groups = "drop") |>
  arrange(desc(total)) |>
  pull(mnemonic)

heat <- heat |>
  mutate(mnemonic = factor(mnemonic, levels = mnem_order))

plot <- ggplot(
  heat,
  aes(x = mnemonic, y = resulttype, fill = fault_rate)
) +
  geom_tile(colour = "white") +
  scale_fill_viridis_c(name = "Fault rate", trans = "log10") +
  labs(
    title = "Instruction / Fault Rate Heatmap (Normalized)",
    x = "Instruction",
    y = "Fault Type"
  ) +
  theme_minimal() +
  theme(
    axis.text.x = element_text(angle = 90, hjust = 1),
    panel.grid = element_blank(),
    plot.title = element_text(size = 14, face = "bold")
  )

suffix <- gsub("^faults|\\.csv$", "", csv_suffix)
outfile <- paste0("injections/instr_fault_rate_heatmap", suffix, ".svg")
ggsave(outfile, plot = plot, width = 12, height = 6)
print(paste("Saved", outfile))