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

# Usage: Rscript combined_instr_fault_correlation.r exp_abspath1 exp_abspath2 ... [faults_file]

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

# TODO: I should probably stop duplicating this each time
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
}

# Don't use counts from faults.csv, as that would correlate completely
# because only mnemonics with faults are listed there
# NOTE: Doesn't match selected filters for faults.csv
mnemonics_file <- "mnemonics.csv"

freq_data <- data.frame()
faults_data <- data.frame()

# TODO: I should probably stop duplicating this each time
for (arg in exp_args) {
  mnem_path <- file.path(arg, mnemonics_file)
  faults_path <- file.path(arg, csv_suffix)

  if (!file.exists(mnem_path)) {
    warning(paste("Missing:", mnem_path))
    next
  }
  if (!file.exists(faults_path)) {
    warning(paste("Missing:", faults_path))
    next
  }

  mdf <- read_csv(mnem_path, col_types = cols())
  mdf$experiment <- basename(arg)
  freq_data <- bind_rows(freq_data, mdf)

  fdf <- read_csv(faults_path, col_types = cols())
  fdf$experiment <- basename(arg)
  faults_data <- bind_rows(faults_data, fdf)
}

if (nrow(freq_data) == 0 || nrow(faults_data) == 0) {
  stop("No data loaded")
}

# x-axis: mnemonics.csv counts
instr_freq <- freq_data |>
  filter(!is.na(mnemonic), mnemonic != "NULL") |>
  group_by(mnemonic) |>
  summarise(instr_count = sum(count, na.rm = TRUE), .groups = "drop")

# y-axis: no OK_MARKERs, sum GROUP1_MARKER + TRAP.
marker_count <- faults_data |>
  filter(!is.na(mnemonic), mnemonic != "NULL") |>
  filter(resulttype != "OK_MARKER") |>
  mutate(
    resulttype = ifelse(resulttype == "GROUP1_MARKER", "TRAP", resulttype)
  ) |>
  group_by(mnemonic) |>
  summarise(marker_count = sum(faults, na.rm = TRUE), .groups = "drop")

correlation <- instr_freq |>
  inner_join(marker_count, by = "mnemonic") |>
  filter(instr_count > 0, marker_count > 0)

if (nrow(correlation) < 2) {
  stop("Not enough mnemonics to compute a correlation")
}

cor_raw <- cor(
  correlation$instr_count,
  correlation$marker_count,
  method = "pearson"
)
cor_log <- cor(
  log10(correlation$instr_count),
  log10(correlation$marker_count),
  method = "pearson"
)
cat(sprintf("Pearson correlation (raw):     %.4f\n", cor_raw))
cat(sprintf("Pearson correlation (log10):   %.4f\n", cor_log))

plot <- ggplot(
  correlation,
  aes(x = instr_count, y = marker_count)
) +
  geom_smooth(
    method = "lm",
    se = FALSE,
    colour = "grey50",
    linetype = "dashed"
  ) +
  geom_point(aes(colour = mnemonic), size = 3, alpha = 0.8) +
  geom_text(
    aes(label = mnemonic),
    size = 3,
    vjust = -0.8,
    check_overlap = TRUE
  ) +
  scale_x_log10(name = "Instruction Executions") +
  scale_y_log10(name = "Fault Count") +
  labs(
    # title = sprintf(
    #   "Instruction / Fault Correlation (r_raw = %.4f, r_log = %.4f)",
    #   cor_raw,
    #   cor_log
    # ),
    title = "Instruction / Fault Correlation",
    colour = "Mnemonic"
  ) +
  theme_minimal() +
  theme(
    legend.position = "none",
    plot.title = element_text(size = 14, face = "bold")
  )

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