server-scripts/qhtcp-workflow/apps/r/calculate_interaction_zscores.R

suppressMessages({
  library("ggplot2")
  library("plotly")
  library("htmlwidgets")
  library("dplyr")
  library("rlang")
  library("ggthemes")
  library("data.table")
  library("unix")
})

options(warn = 2)

# Set the memory limit to 30GB (30 * 1024 * 1024 * 1024 bytes)
soft_limit <- 30 * 1024 * 1024 * 1024
hard_limit <- 30 * 1024 * 1024 * 1024
rlimit_as(soft_limit, hard_limit)

# Constants for configuration
plot_width <- 14
plot_height <- 9
base_size <- 14

parse_arguments <- function() {
  args <- if (interactive()) {
    c(
      "/home/bryan/documents/develop/hartmanlab/qhtcp-workflow/out/20240116_jhartman2_DoxoHLD/20240116_jhartman2_DoxoHLD",
      "/home/bryan/documents/develop/hartmanlab/qhtcp-workflow/apps/r/SGD_features.tab",
      "/home/bryan/documents/develop/hartmanlab/qhtcp-workflow/out/20240116_jhartman2_DoxoHLD/easy/20240116_jhartman2_DoxoHLD/results_std.txt",
      "/home/bryan/documents/develop/hartmanlab/qhtcp-workflow/out/20240116_jhartman2_DoxoHLD/20240822_jhartman2_DoxoHLD/exp1",
      "Experiment 1: Doxo versus HLD",
      3,
      "/home/bryan/documents/develop/hartmanlab/qhtcp-workflow/out/20240116_jhartman2_DoxoHLD/20240822_jhartman2_DoxoHLD/exp2",
      "Experiment 2: HLD versus Doxo",
      3
    )
  } else {
    commandArgs(trailingOnly = TRUE)
  }
  
  # Extract paths, names, and standard deviations
  paths <- args[seq(4, length(args), by = 3)]
  names <- args[seq(5, length(args), by = 3)]
  sds <- as.numeric(args[seq(6, length(args), by = 3)])
  
  # Normalize paths
  normalized_paths <- normalizePath(paths, mustWork = FALSE)
  
  # Create named list of experiments
  experiments <- list()
  for (i in seq_along(paths)) {
    experiments[[names[i]]] <- list(
      path = normalized_paths[i],
      sd = sds[i]
    )
  }
  
  list(
    out_dir = normalizePath(args[1], mustWork = FALSE),
    sgd_gene_list = normalizePath(args[2], mustWork = FALSE),
    easy_results_file = normalizePath(args[3], mustWork = FALSE),
    experiments = experiments
  )
}

args <- parse_arguments()

# Should we keep output in exp dirs or combine in the study output dir?
# dir.create(file.path(args$out_dir, "zscores"), showWarnings = FALSE)
# dir.create(file.path(args$out_dir, "zscores", "qc"), showWarnings = FALSE)

# Define themes and scales
theme_publication <- function(base_size = 14, base_family = "sans", legend_position = "bottom") {
  theme_foundation <- ggplot2::theme_grey(base_size = base_size, base_family = base_family)
  
  theme_foundation %+replace%
    theme(
      plot.title = element_text(face = "bold", size = rel(1.2), hjust = 0.5),
      text = element_text(),
      panel.background = element_rect(colour = NA),
      plot.background = element_rect(colour = NA),
      panel.border = element_rect(colour = NA),
      axis.title = element_text(face = "bold", size = rel(1)),
      axis.title.y = element_text(angle = 90, vjust = 2),
      axis.title.x = element_text(vjust = -0.2),
      axis.line = element_line(colour = "black"),
      panel.grid.major = element_line(colour = "#f0f0f0"),
      panel.grid.minor = element_blank(),
      legend.key = element_rect(colour = NA),
      legend.position = legend_position,
      legend.direction = ifelse(legend_position == "right", "vertical", "horizontal"),
      plot.margin = unit(c(10, 5, 5, 5), "mm"),
      strip.background = element_rect(colour = "#f0f0f0", fill = "#f0f0f0"),
      strip.text = element_text(face = "bold")
    )
}

scale_fill_publication <- function(...) {
  discrete_scale("fill", "Publication", manual_pal(values = c(
    "#386cb0", "#fdb462", "#7fc97f", "#ef3b2c", "#662506",
    "#a6cee3", "#fb9a99", "#984ea3", "#ffff33"
  )), ...)
}

scale_colour_publication <- function(...) {
  discrete_scale("colour", "Publication", manual_pal(values = c(
    "#386cb0", "#fdb462", "#7fc97f", "#ef3b2c", "#662506",
    "#a6cee3", "#fb9a99", "#984ea3", "#ffff33"
  )), ...)
}

# Load the initial dataframe from the easy_results_file
load_and_process_data <- function(easy_results_file, sd = 3) {
  df <- read.delim(easy_results_file, skip = 2, as.is = TRUE, row.names = 1, strip.white = TRUE)
  
  df <- df %>%
    filter(!(.[[1]] %in% c("", "Scan"))) %>%
    filter(!is.na(ORF) & ORF != "" & !Gene %in% c("BLANK", "Blank", "blank") & Drug != "BMH21") %>%
    # Rename columns
    rename(L = l, num = Num., AUC = AUC96, scan = Scan, last_bg = LstBackgrd, first_bg = X1stBackgrd) %>%
    mutate(
      across(c(Col, Row, num, L, K, r, scan, AUC, last_bg, first_bg), as.numeric),
      delta_bg = last_bg - first_bg,
      delta_bg_tolerance = mean(delta_bg, na.rm = TRUE) + (sd * sd(delta_bg, na.rm = TRUE)),
      NG = if_else(L == 0 & !is.na(L), 1, 0),
      DB = if_else(delta_bg >= delta_bg_tolerance, 1, 0),
      SM = 0,
      OrfRep = if_else(ORF == "YDL227C", "YDL227C", OrfRep), # should these be hardcoded?
      conc_num = as.numeric(gsub("[^0-9\\.]", "", Conc)),
      conc_num_factor = as.factor(conc_num)
      # conc_num_factor = factor(conc_num, levels = sort(unique(conc_num)))
      # conc_num_numeric = as.numeric(conc_num_factor) - 1
    )
  
  return(df)
}

# Update Gene names using the SGD gene list
update_gene_names <- function(df, sgd_gene_list) {
  # Load SGD gene list
  genes <- read.delim(file = sgd_gene_list,
    quote = "", header = FALSE,
    colClasses = c(rep("NULL", 3), rep("character", 2), rep("NULL", 11)))
  
  # Create a named vector for mapping ORF to GeneName
  gene_map <- setNames(genes$V5, genes$V4)
  # Vectorized match to find the GeneName from gene_map
  mapped_genes <- gene_map[df$ORF]
  # Replace NAs in mapped_genes with original Gene names (preserves existing Gene names if ORF is not found)
  updated_genes <- ifelse(is.na(mapped_genes) | df$OrfRep == "YDL227C", df$Gene, mapped_genes)
  # Ensure Gene is not left blank or incorrectly updated to "OCT1"
  df <- df %>%
    mutate(Gene = ifelse(updated_genes == "" | updated_genes == "OCT1", OrfRep, updated_genes))
  
  return(df)
}

# Calculate summary statistics for all variables
calculate_summary_stats <- function(df, variables, group_vars = c("OrfRep", "conc_num", "conc_num_factor")) {

  # Summarize the variables within the grouped data
  summary_stats <- df %>%
    group_by(across(all_of(group_vars))) %>%
    summarise(
      N = sum(!is.na(L)),
      across(all_of(variables), list(
        mean = ~mean(., na.rm = TRUE),
        median = ~median(., na.rm = TRUE),
        max = ~ ifelse(all(is.na(.)), NA, max(., na.rm = TRUE)),
        min = ~ ifelse(all(is.na(.)), NA, min(., na.rm = TRUE)),
        sd = ~sd(., na.rm = TRUE),
        se = ~ ifelse(all(is.na(.)), NA, sd(., na.rm = TRUE) / sqrt(sum(!is.na(.)) - 1))
      ), .names = "{.fn}_{.col}")
    )

  # print(summary_stats)

  # Prevent .x and .y suffix issues by renaming columns
  df_cleaned <- df %>%
    select(-any_of(setdiff(names(summary_stats), group_vars))) # Avoid duplicate columns in the final join
  
  # Join the stats back to the original dataframe
  df_with_stats <- left_join(df_cleaned, summary_stats, by = group_vars)

  return(list(summary_stats = summary_stats, df_with_stats = df_with_stats))
}

calculate_interaction_scores <- function(df, max_conc, variables, group_vars = c("OrfRep", "Gene", "num")) {

  # Calculate total concentration variables
  total_conc_num <- length(unique(df$conc_num))

  # Pull the background means and standard deviations from zero concentration
  bg_means <- list(
    L = df %>% filter(conc_num_factor == 0) %>% pull(mean_L) %>% first(),
    K = df %>% filter(conc_num_factor == 0) %>% pull(mean_K) %>% first(),
    r = df %>% filter(conc_num_factor == 0) %>% pull(mean_r) %>% first(),
    AUC = df %>% filter(conc_num_factor == 0) %>% pull(mean_AUC) %>% first()
  )

  bg_sd <- list(
    L = df %>% filter(conc_num_factor == 0) %>% pull(sd_L) %>% first(),
    K = df %>% filter(conc_num_factor == 0) %>% pull(sd_K) %>% first(),
    r = df %>% filter(conc_num_factor == 0) %>% pull(sd_r) %>% first(),
    AUC = df %>% filter(conc_num_factor == 0) %>% pull(sd_AUC) %>% first()
  )

  stats <- df %>%
    group_by(OrfRep, Gene, num, conc_num, conc_num_factor) %>%
    summarise(
      N = sum(!is.na(L)),
      NG = sum(NG, na.rm = TRUE),
      DB = sum(DB, na.rm = TRUE),
      SM = sum(SM, na.rm = TRUE),
      across(all_of(variables), list(
        mean = ~mean(., na.rm = TRUE),
        median = ~median(., na.rm = TRUE),
        max = ~ifelse(all(is.na(.)), NA, max(., na.rm = TRUE)),
        min = ~ifelse(all(is.na(.)), NA, min(., na.rm = TRUE)),
        sd = ~sd(., na.rm = TRUE),
        se = ~ifelse(sum(!is.na(.)) > 1, sd(., na.rm = TRUE) / sqrt(sum(!is.na(.)) - 1), NA)
      ), .names = "{.fn}_{.col}")
    )

  stats <- df %>%
    group_by(OrfRep, Gene, num) %>%
    mutate(
      WT_L = mean_L,
      WT_K = mean_K,
      WT_r = mean_r,
      WT_AUC = mean_AUC,
      WT_sd_L = sd_L,
      WT_sd_K = sd_K,
      WT_sd_r = sd_r,
      WT_sd_AUC = sd_AUC
    )

  stats <- stats %>%
    mutate(
      Raw_Shift_L = first(mean_L) - bg_means$L,
      Raw_Shift_K = first(mean_K) - bg_means$K,
      Raw_Shift_r = first(mean_r) - bg_means$r,
      Raw_Shift_AUC = first(mean_AUC) - bg_means$AUC,
      Z_Shift_L = first(Raw_Shift_L) / bg_sd$L,
      Z_Shift_K = first(Raw_Shift_K) / bg_sd$K,
      Z_Shift_r = first(Raw_Shift_r) / bg_sd$r,
      Z_Shift_AUC = first(Raw_Shift_AUC) / bg_sd$AUC
    )

  stats <- stats %>%
    mutate(
      Exp_L = WT_L + Raw_Shift_L,
      Exp_K = WT_K + Raw_Shift_K,
      Exp_r = WT_r + Raw_Shift_r,
      Exp_AUC = WT_AUC + Raw_Shift_AUC,
      Delta_L = mean_L - Exp_L,
      Delta_K = mean_K - Exp_K,
      Delta_r = mean_r - Exp_r,
      Delta_AUC = mean_AUC - Exp_AUC
    )

  stats <- stats %>%
    mutate(
      Delta_L = if_else(NG == 1, mean_L - WT_L, Delta_L),
      Delta_K = if_else(NG == 1, mean_K - WT_K, Delta_K),
      Delta_r = if_else(NG == 1, mean_r - WT_r, Delta_r),
      Delta_AUC = if_else(NG == 1, mean_AUC - WT_AUC, Delta_AUC),
      Delta_L = if_else(SM == 1, mean_L - WT_L, Delta_L)
    )

  stats <- stats %>%
    mutate(
      Zscore_L = Delta_L / WT_sd_L,
      Zscore_K = Delta_K / WT_sd_K,
      Zscore_r = Delta_r / WT_sd_r,
      Zscore_AUC = Delta_AUC / WT_sd_AUC
    )

  # Calculate linear models
  lm_L <- lm(Delta_L ~ conc_num_factor, data = stats)
  lm_K <- lm(Delta_K ~ conc_num_factor, data = stats)
  lm_r <- lm(Delta_r ~ conc_num_factor, data = stats)
  lm_AUC <- lm(Delta_AUC ~ conc_num_factor, data = stats)

  interactions <- stats %>%
    group_by(OrfRep, Gene, num) %>%
    summarise(
      OrfRep = first(OrfRep),
      Gene = first(Gene),
      num = first(num),
      conc_num = first(conc_num),
      conc_num_factor = first(conc_num_factor),
      Raw_Shift_L = first(Raw_Shift_L),
      Raw_Shift_K = first(Raw_Shift_K),
      Raw_Shift_r = first(Raw_Shift_r),
      Raw_Shift_AUC = first(Raw_Shift_AUC),
      Z_Shift_L = first(Z_Shift_L),
      Z_Shift_K = first(Z_Shift_K),
      Z_Shift_r = first(Z_Shift_r),
      Z_Shift_AUC = first(Z_Shift_AUC),
      Sum_Zscore_L = sum(Zscore_L, na.rm = TRUE),
      Sum_Zscore_K = sum(Zscore_K, na.rm = TRUE),
      Sum_Zscore_r = sum(Zscore_r, na.rm = TRUE),
      Sum_Zscore_AUC = sum(Zscore_AUC, na.rm = TRUE),
      lm_Score_L = max_conc * coef(lm_L)[2] + coef(lm_L)[1],
      lm_Score_K = max_conc * coef(lm_K)[2] + coef(lm_K)[1],
      lm_Score_r = max_conc * coef(lm_r)[2] + coef(lm_r)[1],
      lm_Score_AUC = max_conc * coef(lm_AUC)[2] + coef(lm_AUC)[1],
      R_Squared_L = summary(lm_L)$r.squared,
      R_Squared_K = summary(lm_K)$r.squared,
      R_Squared_r = summary(lm_r)$r.squared,
      R_Squared_AUC = summary(lm_AUC)$r.squared,
      lm_intercept_L = coef(lm_L)[1],
      lm_slope_L = coef(lm_L)[2],
      lm_intercept_K = coef(lm_K)[1],
      lm_slope_K = coef(lm_K)[2],
      lm_intercept_r = coef(lm_r)[1],
      lm_slope_r = coef(lm_r)[2],
      lm_intercept_AUC = coef(lm_AUC)[1],
      lm_slope_AUC = coef(lm_AUC)[2],
      NG = sum(NG, na.rm = TRUE),
      DB = sum(DB, na.rm = TRUE),
      SM = sum(SM, na.rm = TRUE)
    )

  num_non_removed_concs <- total_conc_num - sum(stats$DB, na.rm = TRUE) - 1
  
  interactions <- interactions %>%
    mutate(
      Avg_Zscore_L = Sum_Zscore_L / num_non_removed_concs,
      Avg_Zscore_K = Sum_Zscore_K / num_non_removed_concs,
      Avg_Zscore_r = Sum_Zscore_r / (total_conc_num - 1),
      Avg_Zscore_AUC = Sum_Zscore_AUC / (total_conc_num - 1),
      Z_lm_L = (lm_Score_L - mean(lm_Score_L, na.rm = TRUE)) / sd(lm_Score_L, na.rm = TRUE),
      Z_lm_K = (lm_Score_K - mean(lm_Score_K, na.rm = TRUE)) / sd(lm_Score_K, na.rm = TRUE),
      Z_lm_r = (lm_Score_r - mean(lm_Score_r, na.rm = TRUE)) / sd(lm_Score_r, na.rm = TRUE),
      Z_lm_AUC = (lm_Score_AUC - mean(lm_Score_AUC, na.rm = TRUE)) / sd(lm_Score_AUC, na.rm = TRUE)
    ) %>%
    arrange(desc(Z_lm_L), desc(NG))

  # Declare column order for output
  calculations <- stats %>%
    select(
      "OrfRep", "Gene", "conc_num", "conc_num_factor", "N",
      "mean_L", "mean_K", "mean_r", "mean_AUC",
      "median_L", "median_K", "median_r", "median_AUC",
      "sd_L", "sd_K", "sd_r", "sd_AUC",
      "se_L", "se_K", "se_r", "se_AUC",
      "Raw_Shift_L", "Raw_Shift_K", "Raw_Shift_r", "Raw_Shift_AUC",
      "Z_Shift_L", "Z_Shift_K", "Z_Shift_r", "Z_Shift_AUC",
      "WT_L", "WT_K", "WT_r", "WT_AUC",
      "WT_sd_L", "WT_sd_K", "WT_sd_r", "WT_sd_AUC",
      "Exp_L", "Exp_K", "Exp_r", "Exp_AUC",
      "Delta_L", "Delta_K", "Delta_r", "Delta_AUC",
      "Zscore_L", "Zscore_K", "Zscore_r", "Zscore_AUC",
      "NG", "SM", "DB")
    
  calculations_joined <- df %>% select(-any_of(setdiff(names(calculations), c("OrfRep", "Gene", "num", "conc_num", "conc_num_factor"))))
  calculations_joined <- left_join(calculations_joined, calculations, by = c("OrfRep", "Gene", "num", "conc_num", "conc_num_factor"))

  interactions_joined <- df %>% select(-any_of(setdiff(names(interactions), c("OrfRep", "Gene", "num", "conc_num", "conc_num_factor"))))
  interactions_joined <- left_join(interactions_joined, interactions, by = c("OrfRep", "Gene", "num", "conc_num", "conc_num_factor"))

  return(list(calculations = calculations, interactions = interactions, interactions_joined = interactions_joined,
    calculations_joined = calculations_joined))
}

generate_and_save_plots <- function(output_dir, file_name, plot_configs, grid_layout = NULL) {
  message("Generating ", file_name, ".pdf and ", file_name, ".html")

  # Prepare lists to collect plots
  static_plots <- list()
  plotly_plots <- list()

  for (i in seq_along(plot_configs)) {
    config <- plot_configs[[i]]
    df <- config$df

    # Build the aes_mapping based on config
    aes_mapping <- if (is.null(config$color_var)) {
      if (is.null(config$y_var)) {
        aes(x = .data[[config$x_var]])
      } else {
        aes(x = .data[[config$x_var]], y = .data[[config$y_var]])
      }
    } else {
      if (is.null(config$y_var)) {
        aes(x = .data[[config$x_var]], color = as.factor(.data[[config$color_var]]))
      } else {
        aes(x = .data[[config$x_var]], y = .data[[config$y_var]], color = as.factor(.data[[config$color_var]]))
      }
    }

    # Start building the plot with aes_mapping
    plot_base <- ggplot(df, aes_mapping)

    # Use appropriate helper function based on plot type
    plot <- switch(config$plot_type,
      "scatter" = generate_scatter_plot(plot_base, config),
      "box" = generate_box_plot(plot_base, config),
      "density" = plot_base + geom_density(),
      "bar" = plot_base + geom_bar(),
      plot_base  # default case if no type matches
    )

    # Apply additional settings if provided
    if (!is.null(config$legend_position)) {
      plot <- plot + theme(legend.position = config$legend_position)
    }

    # Add title and labels if provided
    if (!is.null(config$title)) {
      plot <- plot + ggtitle(config$title)
    }
    if (!is.null(config$x_label)) {
      plot <- plot + xlab(config$x_label)
    }
    if (!is.null(config$y_label)) {
      plot <- plot + ylab(config$y_label)
    }

    # Add interactive tooltips for plotly plots
    tooltip_vars <- c("x", "y")  # default tooltip variables
    if (!is.null(config$tooltip_vars)) {
      tooltip_vars <- config$tooltip_vars
    } else {
      # Include default variables based on config
      if (!is.null(config$delta_bg_point) && config$delta_bg_point) {
        tooltip_vars <- c(tooltip_vars, "OrfRep", "Gene", "delta_bg")
      } else if (!is.null(config$gene_point) && config$gene_point) {
        tooltip_vars <- c(tooltip_vars, "OrfRep", "Gene")
      } else {
        # Include x and y variables by default
        tooltip_vars <- c("x", "y")
      }
    }

    # Convert to plotly object
    plotly_plot <- ggplotly(plot, tooltip = tooltip_vars)
    if (!is.null(config$legend_position) && config$legend_position == "bottom") {
      plotly_plot <- plotly_plot %>% layout(legend = list(orientation = "h"))
    }

    # Add plots to lists
    static_plots[[i]] <- plot
    plotly_plots[[i]] <- plotly_plot
  }

  # Save static PDF plots
  pdf(file.path(output_dir, paste0(file_name, ".pdf")), width = 14, height = 9)
  lapply(static_plots, print)
  dev.off()

  # Combine and save interactive HTML plots
  combined_plot <- subplot(plotly_plots, nrows = grid_layout$nrow %||% length(plotly_plots), margin = 0.05)
  saveWidget(combined_plot, file = file.path(output_dir, paste0(file_name, ".html")), selfcontained = TRUE)
}

generate_scatter_plot <- function(plot, config) {
  
  if (!is.null(config$delta_bg_point) && config$delta_bg_point) {
    plot <- plot + geom_point(
      shape = config$shape %||% 3,
      size = config$size %||% 0.2
    )
  } else if (!is.null(config$gene_point) && config$gene_point) {
    plot <- plot + geom_point(
      shape = config$shape %||% 3,
      size = config$size %||% 0.2,
      position = "jitter"
    )
  } else {
    plot <- plot + geom_point(
      shape = config$shape %||% 3,
      position = if (!is.null(config$position) && config$position == "jitter") "jitter" else "identity",
      size = config$size %||% 0.1
    )
  }
  
  # Add smooth line if specified
  if (!is.null(config$add_smooth) && config$add_smooth) {
    if (!is.null(config$lm_line)) {
      plot <- plot + geom_abline(intercept = config$lm_line$intercept, slope = config$lm_line$slope, color = "blue")
    } else {
      plot <- plot + geom_smooth(method = "lm", se = FALSE, color = "blue")
    }
  }
  
  # Add SD bands if specified
  if (!is.null(config$sd_band_values)) {
    for (sd_band in config$sd_band_values) {
      plot <- plot +
        annotate("rect", xmin = -Inf, xmax = Inf, ymin = sd_band, ymax = Inf, fill = "#542788", alpha = 0.3) +
        annotate("rect", xmin = -Inf, xmax = Inf, ymin = -sd_band, ymax = -Inf, fill = "orange", alpha = 0.3) +
        geom_hline(yintercept = c(-sd_band, sd_band), color = "gray")
    }
  }
  
  # Add error bars if specified
  if (!is.null(config$error_bar) && config$error_bar && !is.null(config$y_var)) {
    y_mean_col <- paste0("mean_", config$y_var)
    y_sd_col <- paste0("sd_", config$y_var)
    plot <- plot + geom_errorbar(
      aes(
        ymin = !!sym(y_mean_col) - !!sym(y_sd_col),
        ymax = !!sym(y_mean_col) + !!sym(y_sd_col)
      ),
      alpha = 0.3
    )
  }
  
  # Add x-axis customization if specified
  if (!is.null(config$x_breaks) && !is.null(config$x_labels) && !is.null(config$x_label)) {
    plot <- plot + scale_x_discrete(
      name = config$x_label,
      breaks = config$x_breaks,
      labels = config$x_labels
    )
  }
  
  # Use coord_cartesian for zooming in without removing data outside the range
  if (!is.null(config$coord_cartesian)) {
    plot <- plot + coord_cartesian(ylim = config$coord_cartesian)
  }
  
  # Use scale_y_continuous for setting the y-axis limits
  if (!is.null(config$ylim_vals)) {
    plot <- plot + scale_y_continuous(limits = config$ylim_vals)
  }
  
  # Add annotations if specified
  if (!is.null(config$annotations)) {
    for (annotation in config$annotations) {
      plot <- plot + annotate("text",
        x = annotation$x,
        y = annotation$y,
        label = annotation$label,
        na.rm = TRUE
      )
    }
  }
  
  # Add titles and themes if specified
  if (!is.null(config$title)) {
    plot <- plot + ggtitle(config$title)
  }
  
  if (!is.null(config$legend_position)) {
    plot <- plot + theme(legend.position = config$legend_position)
  }
  
  return(plot)
}


generate_box_plot <- function(plot, config) {
  plot <- plot + geom_boxplot()
  
  if (!is.null(config$x_breaks) && !is.null(config$x_labels) && !is.null(config$x_label)) {
    plot <- plot + scale_x_discrete(
      name = config$x_label,
      breaks = config$x_breaks,
      labels = config$x_labels
    )
  }

  if (!is.null(config$coord_cartesian)) {
    plot <- plot + coord_cartesian(ylim = config$coord_cartesian)
  }

  return(plot)
}

generate_interaction_plot_configs <- function(df, variables) {

  configs <- list()

  limits_map <- list(
    L = c(-65, 65),
    K = c(-65, 65),
    r = c(-0.65, 0.65),
    AUC = c(-6500, 6500)
  )

  df_filtered <- filter_data(df, variables, missing = TRUE, limits_map = limits_map)

  # Define annotation label functions
  generate_annotation_labels <- function(df, var, annotation_name) {
    switch(annotation_name,
      ZShift = paste("ZShift =", round(df[[paste0("Z_Shift_", var)]], 2)),
      lm_ZScore = paste("lm ZScore =", round(df[[paste0("Z_lm_", var)]], 2)),
      NG = paste("NG =", df$NG),
      DB = paste("DB =", df$DB),
      SM = paste("SM =", df$SM),
      NULL  # Default case for unrecognized annotation names
    )
  }

  # Define annotation positions relative to the y-axis range
  calculate_annotation_positions <- function(y_range) {
    y_min <- min(y_range)
    y_max <- max(y_range)
    y_span <- y_max - y_min

    list(
      ZShift = y_max - 0.1 * y_span,
      lm_ZScore = y_max - 0.2 * y_span,
      NG = y_min + 0.2 * y_span,
      DB = y_min + 0.1 * y_span,
      SM = y_min + 0.05 * y_span
    )
  }

  # Create configurations for each variable
  for (variable in variables) {

    y_range <- limits_map[[variable]]
    annotation_positions <- calculate_annotation_positions(y_range)
    lm_line <- list(
      intercept = df_filtered[[paste0("lm_intercept_", variable)]],
      slope = df_filtered[[paste0("lm_slope_", variable)]]
    )

    # Determine x-axis midpoint
    num_levels <- length(levels(df_filtered$conc_num_factor))
    x_pos <- (1 + num_levels) / 2  # Midpoint of x-axis

    # Generate annotations
    annotations <- lapply(names(annotation_positions), function(annotation_name) {
      label <- generate_annotation_labels(df_filtered, variable, annotation_name)
      y_pos <- annotation_positions[[annotation_name]]
      if (!is.null(label)) {
        list(x = x_pos, y = y_pos, label = label)
      } else {
        message(paste("Warning: No annotation found for", annotation_name))
        NULL
      }
    })

    # Remove NULL annotations
    annotations <- Filter(Negate(is.null), annotations)

    # Shared plot settings
    plot_settings <- list(
      df = df_filtered,
      x_var = "conc_num_factor",
      y_var = variable,
      ylim_vals = y_range,
      annotations = annotations,
      lm_line = lm_line,
      x_breaks = levels(df_filtered$conc_num_factor),
      x_labels = levels(df_filtered$conc_num_factor),
      x_label = unique(df_filtered$Drug[1]),
      coord_cartesian = y_range  # Use the actual y-limits
    )

    # Scatter plot config
    configs[[length(configs) + 1]] <- modifyList(plot_settings, list(
      plot_type = "scatter",
      title = sprintf("%s      %s", df_filtered$OrfRep[1], df_filtered$Gene[1]),
      error_bar = TRUE,
      position = "jitter"
    ))

    # Box plot config
    configs[[length(configs) + 1]] <- modifyList(plot_settings, list(
      plot_type = "box",
      title = sprintf("%s      %s (Boxplot)", df_filtered$OrfRep[1], df_filtered$Gene[1]),
      error_bar = FALSE
    ))
  }

  return(configs)
}

generate_rank_plot_configs <- function(df_filtered, is_lm = FALSE, adjust = FALSE) {
  
  # Define SD bands
  sd_bands <- c(1, 2, 3)
  
  # Define variables for Avg ZScore and Rank Avg ZScore plots
  variables <- c("r", "L", "K", "AUC")
  
  # Initialize list to store plot configurations
  configs <- list()
  
  # SD-based plots for L and K
  for (variable in c("L", "K")) {
    for (sd_band in sd_bands) {
      
      # Determine columns based on whether it's lm or not
      if (is_lm) {
        rank_var <- paste0(variable, "_Rank_lm")
        zscore_var <- paste0("Z_lm_", variable)
        y_label <- paste("Int Z score", variable)
      } else {
        rank_var <- paste0(variable, "_Rank")
        zscore_var <- paste0("Avg_Zscore_", variable)
        y_label <- paste("Avg Z score", variable)
      }
      
      # Annotated Plot Configuration
      configs[[length(configs) + 1]] <- list(
        df = df_filtered,
        x_var = rank_var,
        y_var = zscore_var,
        plot_type = "scatter",
        title = paste(y_label, "vs. Rank for", variable, "above", sd_band, "SD"),
        sd_band = sd_band,
        annotations = list(
          list(
            x = median(df_filtered[[rank_var]], na.rm = TRUE),
            y = 10,
            label = paste("Deletion Enhancers =", sum(df_filtered[[zscore_var]] >= sd_band, na.rm = TRUE))
          ),
          list(
            x = median(df_filtered[[rank_var]], na.rm = TRUE),
            y = -10,
            label = paste("Deletion Suppressors =", sum(df_filtered[[zscore_var]] <= -sd_band, na.rm = TRUE))
          )
        ),
        sd_band_values = sd_band,
        shape = 3,
        size = 0.1
      )
      
      # Non-Annotated Plot Configuration
      configs[[length(configs) + 1]] <- list(
        df = df_filtered,
        x_var = rank_var,
        y_var = zscore_var,
        plot_type = "scatter",
        title = paste(y_label, "vs. Rank for", variable, "above", sd_band, "SD No Annotations"),
        sd_band = sd_band,
        annotations = NULL,
        sd_band_values = sd_band,
        shape = 3,
        size = 0.1
      )
    }
  }
  
  # Avg ZScore and Rank Avg ZScore Plots for r, L, K, and AUC
  for (variable in variables) {
    for (plot_type in c("Avg_Zscore_vs_lm", "Rank_Avg_Zscore_vs_lm")) {
      
      # Define x and y variables based on plot type
      if (plot_type == "Avg_Zscore_vs_lm") {
        x_var <- paste0("Avg_Zscore_", variable)
        y_var <- paste0("Z_lm_", variable)
        title_suffix <- paste("Avg Zscore vs lm", variable)
      } else if (plot_type == "Rank_Avg_Zscore_vs_lm") {
        x_var <- paste0(variable, "_Rank")
        y_var <- paste0(variable, "_Rank_lm")
        title_suffix <- paste("Rank Avg Zscore vs lm", variable)
      }
      
      # Determine y-axis label
      if (plot_type == "Avg_Zscore_vs_lm") {
        y_label <- paste("Z lm", variable)
      } else {
        y_label <- paste("Rank lm", variable)
      }
      
      # Determine correlation text (R-squared)
      lm_fit <- lm(df_filtered[[y_var]] ~ df_filtered[[x_var]], data = df_filtered)
      r_squared <- summary(lm_fit)$r.squared
      
      # Plot Configuration
      configs[[length(configs) + 1]] <- list(
        df = df_filtered,
        x_var = x_var,
        y_var = y_var,
        plot_type = "scatter",
        title = title_suffix,
        annotations = list(
          list(
            x = 0,
            y = 0,
            label = paste("R-squared =", round(r_squared, 2))
          )
        ),
        sd_band_values = NULL,  # Not applicable
        shape = 3,
        size = 0.1,
        add_smooth = TRUE,
        lm_line = list(intercept = coef(lm_fit)[1], slope = coef(lm_fit)[2]),
        legend_position = "right"
      )
    }
  }
  
  return(configs)
}

generate_correlation_plot_configs <- function(df) {

  configs <- list()

  variables <- c("r", "L", "K", "AUC")

  for (variable in variables) {
    z_lm_var <- paste0("Z_lm_", variable)
    avg_zscore_var <- paste0("Avg_Zscore_", variable)
    lm_r_squared_col <- paste0("lm_R_squared_", variable)

    configs[[length(configs) + 1]] <- list(
      df = df,
      x_var = avg_zscore_var,
      y_var = z_lm_var,
      plot_type = "scatter",
      title = paste("Avg Zscore vs lm", variable),
      color_var = "Overlap",
      correlation_text = paste("R-squared =", round(df[[lm_r_squared_col]][1], 2)),
      shape = 3,
      geom_smooth = TRUE,
      rect = list(xmin = -2, xmax = 2, ymin = -2, ymax = 2),  # To add the geom_rect layer
      annotate_position = list(x = 0, y = 0),  # Position for the R-squared text
      legend_position = "right"
    )
  }

  return(configs)
}

filter_data <- function(df, variables, nf = FALSE, missing = FALSE, adjust = FALSE,
  rank = FALSE, limits_map = NULL, verbose = TRUE) {
  
  # Precompute Column Names for Efficiency
  avg_zscore_cols <- paste0("Avg_Zscore_", variables)
  z_lm_cols <- paste0("Z_lm_", variables)
  
  # # Optional: Validate that the expected columns exist in the dataframe
  # expected_cols <- c(avg_zscore_cols, z_lm_cols, variables)
  # missing_cols <- setdiff(expected_cols, names(df))
  # if (length(missing_cols) > 0) {
  #   stop("The following expected columns are missing from the dataframe: ",
  #        paste(missing_cols, collapse = ", "))
  # }
  
  # Adjust NAs if 'adjust' is TRUE
  if (adjust) {
    if (verbose) message("Replacing NA with 0.001 for Avg_Zscore_ and Z_lm_ columns.")
    df <- df %>%
      mutate(
        across(all_of(avg_zscore_cols), ~ replace_na(., 0.001)),
        across(all_of(z_lm_cols), ~ replace_na(., 0.001))
      )
  }
  
  # Filter Non-Finite Values if 'nf' is TRUE
  if (nf) {
    if (verbose) message("Filtering non-finite values for variables: ", paste(variables, collapse = ", "))
    
    # Identify non-finite rows for logging
    non_finite_df <- df %>%
      filter(if_any(all_of(variables), ~ !is.finite(.)))
    
    if (verbose && nrow(non_finite_df) > 0) {
      message("Non-finite rows for variables ", paste(variables, collapse = ", "), ":")
      print(non_finite_df)
    }
    
    # Keep only rows where all specified variables are finite
    df <- df %>%
      filter(if_all(all_of(variables), ~ is.finite(.)))
  }
  
  # Filter Missing Values if 'missing' is TRUE
  if (missing) {
    if (verbose) message("Filtering missing values for variables: ", paste(variables, collapse = ", "))
    
    # Identify missing rows for logging
    missing_df <- df %>%
      filter(if_any(all_of(variables), ~ is.na(.)))
    
    if (verbose && nrow(missing_df) > 0) {
      message("Missing data for variables ", paste(variables, collapse = ", "), ":")
      print(missing_df)
    }
    
    # Keep only rows where all specified variables are not missing
    df <- df %>%
      filter(if_all(all_of(variables), ~ !is.na(.)))
  }
  
  # Apply Limits from 'limits_map' if Provided
  if (!is.null(limits_map)) {
    for (variable in names(limits_map)) {
      if (variable %in% variables) {
        ylim_vals <- limits_map[[variable]]
        
        if (verbose) message("Applying limits for variable ", variable, ": [", ylim_vals[1], ", ", ylim_vals[2], "].")
        
        # Identify out-of-range data for logging
        out_of_range_df <- df %>%
          filter(.data[[variable]] < ylim_vals[1] | .data[[variable]] > ylim_vals[2])
        
        if (verbose && nrow(out_of_range_df) > 0) {
          message("Out-of-range data for variable ", variable, ":")
          print(out_of_range_df)
        }
        
        # Keep only rows within the specified limits
        df <- df %>%
          filter(.data[[variable]] >= ylim_vals[1] & .data[[variable]] <= ylim_vals[2])
      }
    }
  }
  
  # Calculate Rank Columns if 'rank' is TRUE
  if (rank) {
    if (verbose) message("Calculating rank columns for variables: ", paste(variables, collapse = ", "))
    
    # Create Rank and Rank_lm columns using mutate and across
    df <- df %>%
      mutate(
        # Rank based on Avg_Zscore_
        across(all_of(avg_zscore_cols), ~ rank(., na.last = "keep"), .names = "Rank_Avg_Zscore_{.col}"),
        # Rank_lm based on Z_lm_
        across(all_of(z_lm_cols), ~ rank(., na.last = "keep"), .names = "Rank_lm_Z_lm_{.col}")
      )
    
    # Prepare a named vector for renaming columns: new_name = old_name
    rename_vector <- c(
      setNames(paste0("Rank_", variables), paste0("Rank_Avg_Zscore_", avg_zscore_cols)),
      setNames(paste0("Rank_lm_", variables), paste0("Rank_lm_Z_lm_", z_lm_cols))
    )
    
    # Rename the rank columns in a single step
    df <- df %>%
      rename(!!!rename_vector)
  }
  
  return(df)
}

main <- function() {
  lapply(names(args$experiments), function(exp_name) {
    exp <- args$experiments[[exp_name]]
    exp_path <- exp$path
    exp_sd <- exp$sd
    out_dir <- file.path(exp_path, "zscores")
    out_dir_qc <- file.path(exp_path, "zscores", "qc")
    dir.create(out_dir, recursive = TRUE, showWarnings = FALSE)
    dir.create(out_dir_qc, recursive = TRUE, showWarnings = FALSE)

    summary_vars <- c("L", "K", "r", "AUC", "delta_bg") # fields to filter and calculate summary stats across
    group_vars <- c("OrfRep", "conc_num", "conc_num_factor") # default fields to group by
    orf_group_vars <- c("OrfRep", "Gene", "num")
    print_vars <- c("OrfRep", "Plate", "scan", "Col", "Row", "num", "OrfRep", "conc_num", "conc_num_factor",
      "delta_bg_tolerance", "delta_bg", "Gene", "L", "K", "r", "AUC", "NG", "DB")
    
    message("Loading and filtering data for experiment: ", exp_name)
    df <- load_and_process_data(args$easy_results_file, sd = exp_sd) %>%
      update_gene_names(args$sgd_gene_list) %>%
      as_tibble()

    # Filter rows above delta background tolerance
    df_above_tolerance <- df %>% filter(DB == 1)
    df_na <- df %>% mutate(across(all_of(summary_vars), ~ ifelse(DB == 1, NA, .)))
    df_no_zeros <- df_na %>% filter(L > 0)
    
    # Save some constants
    max_conc <- max(df$conc_num)
    l_half_median <- (median(df_above_tolerance$L, na.rm = TRUE)) / 2
    k_half_median <- (median(df_above_tolerance$K, na.rm = TRUE)) / 2

    message("Calculating summary statistics before quality control")
    ss <- calculate_summary_stats(df, summary_vars, group_vars = group_vars)
    df_stats <- ss$df_with_stats
    message("Filtering non-finite data")
    df_filtered_stats <- filter_data(df_stats, c("L"), nf = TRUE)

    message("Calculating summary statistics after quality control")
    ss <- calculate_summary_stats(df_na, summary_vars, group_vars = group_vars)
    df_na_ss <- ss$summary_stats
    df_na_stats <- ss$df_with_stats
    write.csv(df_na_ss, file = file.path(out_dir, "summary_stats_all_strains.csv"), row.names = FALSE)
    df_na_filtered_stats <- filter_data(df_na_stats, c("L"), nf = TRUE)

    message("Calculating summary statistics after quality control excluding zero values")
    ss <- calculate_summary_stats(df_no_zeros, summary_vars, group_vars = group_vars)
    df_no_zeros_stats <- ss$df_with_stats
    df_no_zeros_filtered_stats <- filter_data(df_no_zeros_stats, c("L"), nf = TRUE)

    message("Filtering by 2SD of K")
    df_na_within_2sd_k <- df_na_stats %>%
      filter(K >= (mean_K - 2 * sd_K) & K <= (mean_K + 2 * sd_K))
    df_na_outside_2sd_k <- df_na_stats %>%
      filter(K < (mean_K - 2 * sd_K) | K > (mean_K + 2 * sd_K))

    message("Calculating summary statistics for L within 2SD of K")
    # TODO We're omitting the original z_max calculation, not sure if needed?
    ss <- calculate_summary_stats(df_na_within_2sd_k, "L", group_vars = c("conc_num", "conc_num_factor"))
    # df_na_l_within_2sd_k_stats <- ss$df_with_stats
    write.csv(ss$summary_stats, file = file.path(out_dir_qc, "max_observed_L_vals_for_spots_within_2sd_K.csv"), row.names = FALSE)
    
    message("Calculating summary statistics for L outside 2SD of K")
    ss <- calculate_summary_stats(df_na_outside_2sd_k, "L", group_vars = c("conc_num", "conc_num_factor"))
    df_na_l_outside_2sd_k_stats <- ss$df_with_stats
    write.csv(ss$summary_stats, file = file.path(out_dir, "max_observed_L_vals_for_spots_outside_2sd_K.csv"), row.names = FALSE)
    
    # Each plots list corresponds to a file
    l_vs_k_plots <- list(
      list(
        df = df,
        x_var = "L",
        y_var = "K",
        plot_type = "scatter",
        delta_bg_point = TRUE,
        title = "Raw L vs K before quality control",
        color_var = "conc_num_factor",
        error_bar = FALSE,
        legend_position = "right"
      )
    )

    frequency_delta_bg_plots <- list(
      list(
        df = df_filtered_stats,
        x_var = "delta_bg",
        y_var = NULL,
        plot_type = "density",
        title = "Plate analysis by Drug Conc for Delta Background before quality control",
        color_var = "conc_num_factor",
        x_label = "Delta Background",
        y_label = "Density",
        error_bar = FALSE,
        legend_position = "right"),
      list(
        df = df_filtered_stats,
        x_var = "delta_bg",
        y_var = NULL,
        plot_type = "bar",
        title = "Plate analysis by Drug Conc for Delta Background before quality control",
        color_var = "conc_num_factor",
        x_label = "Delta Background",
        y_label = "Count",
        error_bar = FALSE,
        legend_position = "right")
    )

    above_threshold_plots <- list(
      list(
        df = df_above_tolerance,
        x_var = "L",
        y_var = "K",
        plot_type = "scatter",
        delta_bg_point = TRUE,
        title = paste("Raw L vs K for strains above Delta Background threshold of",
          df_above_tolerance$delta_bg_tolerance[[1]], "or above"),
        color_var = "conc_num_factor",
        position = "jitter",
        annotations = list(
          list(
            x = l_half_median,
            y = k_half_median,
            label = paste("# strains above Delta Background tolerance =", nrow(df_above_tolerance))
          )
        ),
        error_bar = FALSE,
        legend_position = "right"
      )
    )

    plate_analysis_plots <- list()
    for (var in summary_vars) {
      for (stage in c("before", "after")) {
        if (stage == "before") {
          df_plot <- df_filtered_stats
        } else {
          df_plot <- df_na_filtered_stats
        }
        
        config <- list(
          df = df_plot,
          x_var = "scan",
          y_var = var,
          plot_type = "scatter",
          title = paste("Plate analysis by Drug Conc for", var, stage, "quality control"),
          error_bar = TRUE,
          color_var = "conc_num_factor",
          position = "jitter")

        plate_analysis_plots <- append(plate_analysis_plots, list(config))
      }
    }

    plate_analysis_boxplots <- list()
    for (var in summary_vars) {
      for (stage in c("before", "after")) {
        if (stage == "before") {
          df_plot <- df_filtered_stats
        } else {
          df_plot <- df_na_filtered_stats
        }
        
        config <- list(
          df = df_plot,
          x_var = "scan",
          y_var = var,
          plot_type = "box",
          title = paste("Plate analysis by Drug Conc for", var, stage, "quality control"),
          error_bar = FALSE,
          color_var = "conc_num_factor")

        plate_analysis_boxplots <- append(plate_analysis_boxplots, list(config))
      }
    }

    plate_analysis_no_zeros_plots <- list()
    for (var in summary_vars) {
      config <- list(
        df = df_no_zeros_filtered_stats,
        x_var = "scan",
        y_var = var,
        plot_type = "scatter",
        title = paste("Plate analysis by Drug Conc for", var, "after quality control"),
        error_bar = TRUE,
        color_var = "conc_num_factor",
        position = "jitter")

      plate_analysis_no_zeros_plots <- append(plate_analysis_no_zeros_plots, list(config))
    }

    plate_analysis_no_zeros_boxplots <- list()
    for (var in summary_vars) {
      config <- list(
        df = df_no_zeros_filtered_stats,
        x_var = "scan",
        y_var = var,
        plot_type = "box",
        title = paste("Plate analysis by Drug Conc for", var, "after quality control"),
        error_bar = FALSE,
        color_var = "conc_num_factor"
      )
      plate_analysis_no_zeros_boxplots <- append(plate_analysis_no_zeros_boxplots, list(config))
    }

    l_outside_2sd_k_plots <- list(
      list(
        df = df_na_l_outside_2sd_k_stats,
        x_var = "L",
        y_var = "K",
        plot_type = "scatter",
        delta_bg_point = TRUE,
        title = "Raw L vs K for strains falling outside 2SD of the K mean at each Conc",
        color_var = "conc_num_factor",
        position = "jitter",
        legend_position = "right"
      )
    )

    delta_bg_outside_2sd_k_plots <- list(
      list(
        df = df_na_l_outside_2sd_k_stats,
        x_var = "delta_bg",
        y_var = "K",
        plot_type = "scatter",
        gene_point = TRUE,
        title = "Delta Background vs K for strains falling outside 2SD of the K mean at each Conc",
        color_var = "conc_num_factor",
        position = "jitter",
        legend_position = "right"
      )
    )

    message("Generating quality control plots")
    generate_and_save_plots(out_dir_qc, "L_vs_K_before_quality_control", l_vs_k_plots)
    generate_and_save_plots(out_dir_qc, "frequency_delta_background", frequency_delta_bg_plots)
    generate_and_save_plots(out_dir_qc, "L_vs_K_above_threshold", above_threshold_plots)
    generate_and_save_plots(out_dir_qc, "plate_analysis", plate_analysis_plots)
    generate_and_save_plots(out_dir_qc, "plate_analysis_boxplots", plate_analysis_boxplots)
    generate_and_save_plots(out_dir_qc, "plate_analysis_no_zeros", plate_analysis_no_zeros_plots)
    generate_and_save_plots(out_dir_qc, "plate_analysis_no_zeros_boxplots", plate_analysis_no_zeros_boxplots)
    generate_and_save_plots(out_dir_qc, "L_vs_K_for_strains_2SD_outside_mean_K", l_outside_2sd_k_plots)
    generate_and_save_plots(out_dir_qc, "delta_background_vs_K_for_strains_2sd_outside_mean_K", delta_bg_outside_2sd_k_plots)

    # TODO: Originally this filtered L NA's
    # Let's try to avoid for now since stats have already been calculated

    # Process background strains
    bg_strains <- c("YDL227C")
    lapply(bg_strains, function(strain) {
      
      message("Processing background strain: ", strain)
      
      # Handle missing data by setting zero values to NA
      # and then removing any rows with NA in L col
      df_bg <- df_na %>%
        filter(OrfRep == strain) %>%
        mutate(
          L = if_else(L == 0, NA, L),
          K = if_else(K == 0, NA, K),
          r = if_else(r == 0, NA, r),
          AUC = if_else(AUC == 0, NA, AUC)
        ) %>%
        filter(!is.na(L))
      
      # Recalculate summary statistics for the background strain
      message("Calculating summary statistics for background strain")
      ss_bg <- calculate_summary_stats(df_bg, summary_vars, group_vars = group_vars)
      summary_stats_bg <- ss_bg$summary_stats
      # df_bg_stats <- ss_bg$df_with_stats
      write.csv(summary_stats_bg,
        file = file.path(out_dir, paste0("SummaryStats_BackgroundStrains_", strain, ".csv")),
        row.names = FALSE)
      
      # Filter reference and deletion strains
      # Formerly X2_RF (reference strains)
      df_reference <- df_na_stats %>%
        filter(OrfRep == strain) %>%
        mutate(SM = 0)
      
      # Formerly X2 (deletion strains)
      df_deletion <- df_na_stats %>%
        filter(OrfRep != strain) %>%
        mutate(SM = 0)

      # Set the missing values to the highest theoretical value at each drug conc for L
      # Leave other values as 0 for the max/min
      reference_strain <- df_reference %>%
        group_by(conc_num) %>%
        mutate(
          max_l_theoretical = max(max_L, na.rm = TRUE),
          L = ifelse(L == 0 & !is.na(L) & conc_num > 0, max_l_theoretical, L),
          SM = ifelse(L >= max_l_theoretical & !is.na(L) & conc_num > 0, 1, SM),
          L = ifelse(L >= max_l_theoretical & !is.na(L) & conc_num > 0, max_l_theoretical, L)) %>%
        ungroup()

      # Ditto for deletion strains
      deletion_strains <- df_deletion %>%
        group_by(conc_num) %>%
        mutate(
          max_l_theoretical = max(max_L, na.rm = TRUE),
          L = ifelse(L == 0 & !is.na(L) & conc_num > 0, max_l_theoretical, L),
          SM = ifelse(L >= max_l_theoretical & !is.na(L) & conc_num > 0, 1, SM),
          L = ifelse(L >= max_l_theoretical & !is.na(L) & conc_num > 0, max_l_theoretical, L)) %>%
        ungroup()

      message("Calculating interaction scores")
      interaction_vars <- c("L", "K", "r", "AUC")

      message("Calculating reference strain(s)")
      reference_results <- calculate_interaction_scores(reference_strain, max_conc, interaction_vars, group_vars = orf_group_vars)
      zscores_calculations_reference <- reference_results$calculations
      zscores_interactions_reference <- reference_results$interactions
      zscores_calculations_reference_joined <- reference_results$calculations_joined
      zscores_interactions_reference_joined <- reference_results$interactions_joined

      message("Calculating deletion strain(s)")
      deletion_results <- calculate_interaction_scores(deletion_strains, max_conc, interaction_vars, group_vars = orf_group_vars)
      zscores_calculations <- deletion_results$calculations
      zscores_interactions <- deletion_results$interactions
      zscores_calculations_joined <- deletion_results$calculations_joined
      zscores_interactions_joined <- deletion_results$interactions_joined

      # Writing Z-Scores to file
      write.csv(zscores_calculations_reference, file = file.path(out_dir, "RF_ZScores_Calculations.csv"), row.names = FALSE)
      write.csv(zscores_calculations, file = file.path(out_dir, "ZScores_Calculations.csv"), row.names = FALSE)
      write.csv(zscores_interactions_reference, file = file.path(out_dir, "RF_ZScores_Interaction.csv"), row.names = FALSE)
      write.csv(zscores_interactions, file = file.path(out_dir, "ZScores_Interaction.csv"), row.names = FALSE)

      # Create interaction plots
      message("Generating reference interaction plots")
      reference_plot_configs <- generate_interaction_plot_configs(zscores_interactions_reference_joined, interaction_vars)
      generate_and_save_plots(out_dir, "RF_interactionPlots", reference_plot_configs, grid_layout = list(ncol = 4, nrow = 3))

      message("Generating deletion interaction plots")
      deletion_plot_configs <- generate_interaction_plot_configs(zscores_interactions_joined, interaction_vars)
      generate_and_save_plots(out_dir, "InteractionPlots", deletion_plot_configs, grid_layout = list(ncol = 4, nrow = 3))

      # Define conditions for enhancers and suppressors
      # TODO Add to study config file?
      threshold <- 2
      enhancer_condition_L <- zscores_interactions$Avg_Zscore_L >= threshold
      suppressor_condition_L <- zscores_interactions$Avg_Zscore_L <= -threshold
      enhancer_condition_K <- zscores_interactions$Avg_Zscore_K >= threshold
      suppressor_condition_K <- zscores_interactions$Avg_Zscore_K <= -threshold
      
      # Subset data
      enhancers_L <- zscores_interactions[enhancer_condition_L, ]
      suppressors_L <- zscores_interactions[suppressor_condition_L, ]
      enhancers_K <- zscores_interactions[enhancer_condition_K, ]
      suppressors_K <- zscores_interactions[suppressor_condition_K, ]
      
      # Save enhancers and suppressors
      message("Writing enhancer/suppressor csv files")
      write.csv(enhancers_L, file = file.path(out_dir, "ZScores_Interaction_Deletion_Enhancers_L.csv"), row.names = FALSE)
      write.csv(suppressors_L, file = file.path(out_dir, "ZScores_Interaction_Deletion_Suppressors_L.csv"), row.names = FALSE)
      write.csv(enhancers_K, file = file.path(out_dir, "ZScores_Interaction_Deletion_Enhancers_K.csv"), row.names = FALSE)
      write.csv(suppressors_K, file = file.path(out_dir, "ZScores_Interaction_Deletion_Suppressors_K.csv"), row.names = FALSE)
      
      # Combine conditions for enhancers and suppressors
      enhancers_and_suppressors_L <- zscores_interactions[enhancer_condition_L | suppressor_condition_L, ]
      enhancers_and_suppressors_K <- zscores_interactions[enhancer_condition_K | suppressor_condition_K, ]
      
      # Save combined enhancers and suppressors
      write.csv(enhancers_and_suppressors_L,
        file = file.path(out_dir, "ZScores_Interaction_Deletion_Enhancers_and_Suppressors_L.csv"), row.names = FALSE)
      write.csv(enhancers_and_suppressors_K,
        file = file.path(out_dir, "ZScores_Interaction_Deletion_Enhancers_and_Suppressors_K.csv"), row.names = FALSE)
      
      # Handle linear model based enhancers and suppressors
      lm_threshold <- 2
      enhancers_lm_L <- zscores_interactions[zscores_interactions$Z_lm_L >= lm_threshold, ]
      suppressors_lm_L <- zscores_interactions[zscores_interactions$Z_lm_L <= -lm_threshold, ]
      enhancers_lm_K <- zscores_interactions[zscores_interactions$Z_lm_K >= lm_threshold, ]
      suppressors_lm_K <- zscores_interactions[zscores_interactions$Z_lm_K <= -lm_threshold, ]
      
      # Save linear model based enhancers and suppressors
      message("Writing linear model enhancer/suppressor csv files")
      write.csv(enhancers_lm_L,
        file = file.path(out_dir, "ZScores_Interaction_Deletion_Enhancers_L_lm.csv"), row.names = FALSE)
      write.csv(suppressors_lm_L,
        file = file.path(out_dir, "ZScores_Interaction_Deletion_Suppressors_L_lm.csv"), row.names = FALSE)
      write.csv(enhancers_lm_K,
        file = file.path(out_dir, "ZScores_Interaction_Deletion_Enhancers_K_lm.csv"), row.names = FALSE)
      write.csv(suppressors_lm_K,
        file = file.path(out_dir, "ZScores_Interaction_Deletion_Suppressors_K_lm.csv"), row.names = FALSE)

      message("Generating rank plots")
      zscores_interactions_joined_filtered <- filter_data(
        zscores_interactions_joined,
        variables,
        missing = TRUE,
        adjust = TRUE,
        rank = TRUE)
      rank_plot_configs <- generate_rank_plot_configs(
        df = zscores_interactions_joined_filtered,
        is_lm = FALSE,
        adjust = TRUE
      )
      generate_and_save_plots(output_dir = out_dir, file_name = "RankPlots",
        plot_configs = rank_plot_configs, grid_layout = list(ncol = 3, nrow = 2))

      message("Generating ranked linear model plots")
      rank_lm_plot_configs <- generate_rank_plot_configs(
        df = zscores_interactions_joined_filtered,
        is_lm = TRUE,
        adjust = TRUE
      )
      generate_and_save_plots(output_dir = out_dir, file_name = "RankPlots_lm",
        plot_configs = rank_lm_plot_configs, grid_layout = list(ncol = 3, nrow = 2))

      message("Filtering and reranking plots")
      # Formerly X_NArm
      zscores_interactions_filtered <- zscores_interactions_joined %>%
        group_by(across(all_of(orf_group_vars))) %>%
        filter(!is.na(Z_lm_L) | !is.na(Avg_Zscore_L)) %>%
        ungroup() %>%
        rowwise() %>%
        mutate(
          lm_R_squared_L = if (n() > 1) summary(lm(Z_lm_L ~ Avg_Zscore_L))$r.squared else NA,
          lm_R_squared_K = if (n() > 1) summary(lm(Z_lm_K ~ Avg_Zscore_K))$r.squared else NA,
          lm_R_squared_r = if (n() > 1) summary(lm(Z_lm_r ~ Avg_Zscore_r))$r.squared else NA,
          lm_R_squared_AUC = if (n() > 1) summary(lm(Z_lm_AUC ~ Avg_Zscore_AUC))$r.squared else NA,

          Overlap = case_when(
            Z_lm_L >= 2 & Avg_Zscore_L >= 2 ~ "Deletion Enhancer Both",
            Z_lm_L <= -2 & Avg_Zscore_L <= -2 ~ "Deletion Suppressor Both",
            Z_lm_L >= 2 & Avg_Zscore_L <= 2 ~ "Deletion Enhancer lm only",
            Z_lm_L <= -2 & Avg_Zscore_L >= -2 ~ "Deletion Suppressor lm only",
            Z_lm_L >= 2 & Avg_Zscore_L <= -2 ~ "Deletion Enhancer lm, Deletion Suppressor Avg Z score",
            Z_lm_L <= -2 & Avg_Zscore_L >= 2 ~ "Deletion Suppressor lm, Deletion Enhancer Avg Z score",
            TRUE ~ "No Effect"
          )
        ) %>%
        ungroup()

      # Re-rank
      zscores_interactions_filtered <- filter_data(
        df = zscores_interactions_filtered,
        variables = interaction_vars,
        missing = TRUE, # TODO what I'm currently having issues with
        rank = TRUE
      )

      rank_plot_filtered_configs <- generate_rank_plot_configs(
        df = zscores_interactions_filtered,
        is_lm = FALSE,
        adjust = FALSE
      )

      message("Generating filtered ranked plots")
      generate_and_save_plots(
        output_dir = out_dir,
        file_name = "RankPlots_na_rm",
        plot_configs = rank_plot_filtered_configs,
        grid_layout = list(ncol = 3, nrow = 2))

      message("Generating filtered ranked linear model plots")
      rank_plot_lm_filtered_configs <- generate_rank_plot_configs(
        df = zscores_interactions_filtered,
        is_lm = TRUE,
        adjust = FALSE
      )
      generate_and_save_plots(
        output_dir = out_dir,
        file_name = "RankPlots_lm_na_rm",
        plot_configs = rank_plot_lm_filtered_configs,
        grid_layout = list(ncol = 3, nrow = 2))

      message("Generating correlation plots")
      correlation_plot_configs <- generate_correlation_plot_configs(zscores_interactions_filtered)
      generate_and_save_plots(
        output_dir = out_dir,
        file_name = "Avg_Zscore_vs_lm_NA_rm",
        plot_configs = correlation_plot_configs,
        grid_layout = list(ncol = 2, nrow = 2))
    })
  })
}
main()