plot_orf_method_comparison.R

#!/usr/bin/env Rscript
#==============================================================================
# Visual comparison: six-frame best ORF vs TransDecoder single-best modes
# (complete-only vs with partials) using orf_comparison_summary.tsv from each
# compare_orf_prediction run.
#
# Usage (local-friendly):
#   cd path/to/orf_method_comparison/SAMPLE_td_singlebest
#   Rscript plot_orf_method_comparison.R .
#
#   Rscript plot_orf_method_comparison.R path/to/orf_method_comparison/SAMPLE_td_singlebest
#
#   # If you only pass the parent orf_method_comparison folder and it contains
#   # exactly one *_td_singlebest run, that folder is used automatically:
#   Rscript plot_orf_method_comparison.R path/to/orf_method_comparison
#
# Expects:
#   <base>/complete/orf_comparison_summary.tsv
#   <base>/with_partials/orf_comparison_summary.tsv
#
# Optional second path / env for MetaEuk compare output (four bars total):
#   Rscript plot_orf_method_comparison.R SAMPLE_td_singlebest SAMPLE_metaeuk
#   Rscript plot_orf_method_comparison.R SAMPLE_td_singlebest /path/to/orf_comparison_summary.tsv
#   export ORF_SUMMARY_METAEUK=/path/to/MMA_ST106_L004_metaeuk/orf_comparison_summary.tsv
# If omitted, looks for sibling folder <stem>_metaeuk/orf_comparison_summary.tsv next to *_td_singlebest.
#
# Optional env: ORF_PLOT_OUT (output prefix; default: <base>/orf_method_comparison_plot)
#==============================================================================
suppressPackageStartupMessages({
  library(readr)
  library(dplyr)
  library(tidyr)
  library(ggplot2)
  library(stringr)
})

read_summary_kv <- function(path) {
  if (!file.exists(path)) {
    stop("Missing file: ", path)
  }
  x <- read_tsv(path, col_names = c("key", "value"), show_col_types = FALSE)
  x$value_num <- suppressWarnings(as.numeric(x$value))
  vals <- x$value_num
  names(vals) <- x$key
  vals
}

getn <- function(v, k, default = NA_real_) {
  if (!k %in% names(v)) {
    return(default)
  }
  out <- v[[k]]
  if (length(out) != 1L || is.na(out)) {
    return(default)
  }
  out
}

has_td_layout <- function(d) {
  file.exists(file.path(d, "complete", "orf_comparison_summary.tsv")) &&
    file.exists(file.path(d, "with_partials", "orf_comparison_summary.tsv"))
}

infer_metaeuk_summary_path <- function(td_base_dir) {
  parent <- dirname(td_base_dir)
  stem <- sub("_td_singlebest$", "", basename(td_base_dir))
  f <- file.path(parent, paste0(stem, "_metaeuk"), "orf_comparison_summary.tsv")
  if (file.exists(f)) {
    return(f)
  }
  NA_character_
}

extract_sample_ids <- function(path_like) {
  b <- basename(path_like)
  out <- character(0)
  m1 <- str_match(b, "(ST\\d+_L\\d+)")
  if (!is.na(m1[1, 2])) {
    out <- c(out, m1[1, 2])
    m1n <- str_match(m1[1, 2], "ST(\\d+)_L(\\d+)")
    if (!is.na(m1n[1, 2]) && !is.na(m1n[1, 3])) {
      out <- c(out, paste0(m1n[1, 2], "_L", m1n[1, 3]))
    }
  }
  m2 <- str_match(b, "MMA_(\\d+)_L(\\d+)")
  if (!is.na(m2[1, 2]) && !is.na(m2[1, 3])) {
    out <- c(out, paste0("ST", m2[1, 2], "_L", m2[1, 3]))
    out <- c(out, paste0(m2[1, 2], "_L", m2[1, 3]))
  }
  unique(out)
}

find_optional_summary <- function(base_dir, method_suffix, args = character(0), env_var = "") {
  # Priority: explicit CLI arg > explicit env var > sibling folder inference.
  if (length(args) >= 1L && nzchar(args[[1]])) {
    p <- normalizePath(args[[1]], winslash = "/", mustWork = TRUE)
    if (dir.exists(p)) {
      f <- file.path(p, "orf_comparison_summary.tsv")
      if (!file.exists(f)) {
        stop("No orf_comparison_summary.tsv in directory: ", p)
      }
      return(f)
    }
    if (!file.exists(p)) {
      stop("Summary path not found: ", p)
    }
    return(p)
  }

  if (nzchar(env_var)) {
    ev <- Sys.getenv(env_var, unset = "")
    if (nzchar(ev)) {
      p <- normalizePath(ev, winslash = "/", mustWork = TRUE)
      if (dir.exists(p)) {
        f <- file.path(p, "orf_comparison_summary.tsv")
        if (!file.exists(f)) {
          stop("No orf_comparison_summary.tsv in: ", p)
        }
        return(f)
      }
      if (!file.exists(p)) {
        stop(env_var, " not found: ", p)
      }
      return(p)
    }
  }

  parent <- dirname(base_dir)
  sample_ids <- extract_sample_ids(base_dir)
  if (length(sample_ids) == 0L) {
    return(NA_character_)
  }
  candidates <- unique(unlist(lapply(sample_ids, function(sid) {
    c(
      file.path(parent, paste0("MMA_", sid, "_", method_suffix), "orf_comparison_summary.tsv"),
      file.path(parent, paste0(sid, "_", method_suffix), "orf_comparison_summary.tsv")
    )
  })))
  hit <- candidates[file.exists(candidates)]
  if (length(hit) >= 1L) {
    return(hit[[1]])
  }
  NA_character_
}

resolve_metaeuk_summary <- function(td_base_dir, args) {
  if (length(args) >= 2L) {
    raw <- args[[2]]
    p <- normalizePath(raw, winslash = "/", mustWork = TRUE)
    if (dir.exists(p)) {
      f <- file.path(p, "orf_comparison_summary.tsv")
      if (!file.exists(f)) {
        stop("No orf_comparison_summary.tsv in directory: ", p)
      }
      return(f)
    }
    if (!file.exists(p)) {
      stop("MetaEuk summary path not found: ", p)
    }
    return(p)
  }
  ev <- Sys.getenv("ORF_SUMMARY_METAEUK", unset = "")
  if (nzchar(ev)) {
    p <- normalizePath(ev, winslash = "/", mustWork = TRUE)
    if (dir.exists(p)) {
      f <- file.path(p, "orf_comparison_summary.tsv")
      if (!file.exists(f)) {
        stop("No orf_comparison_summary.tsv in: ", p)
      }
      return(f)
    }
    if (!file.exists(p)) {
      stop("ORF_SUMMARY_METAEUK not found: ", p)
    }
    return(p)
  }
  infer_metaeuk_summary_path(td_base_dir)
}

find_single_td_singlebest <- function(parent) {
  if (!dir.exists(parent)) {
    return(character(0))
  }
  subs <- file.path(parent, list.files(parent, all.files = FALSE, no.. = TRUE))
  isd <- file.info(subs)$isdir
  subs <- subs[!is.na(isd) & isd]
  hits <- subs[grepl("_td_singlebest$", basename(subs))]
  hits[!is.na(hits)]
}

resolve_base_dir <- function(args) {
  if (length(args) == 0L) {
    cand <- getwd()
    if (has_td_layout(cand)) {
      return(normalizePath(cand, winslash = "/", mustWork = TRUE))
    }
    alt <- find_single_td_singlebest(cand)
    if (length(alt) == 1L) {
      message("Using: ", alt[[1]])
      return(normalizePath(alt[[1]], winslash = "/", mustWork = TRUE))
    }
    stop(
      "No arguments: cwd is not a SAMPLE_td_singlebest folder, ",
      "and does not contain exactly one *_td_singlebest subdir.\n",
      "Usage:\n",
      "  cd SAMPLE_td_singlebest && Rscript plot_orf_method_comparison.R\n",
      "  cd orf_method_comparison && Rscript plot_orf_method_comparison.R   # one *_td_singlebest only\n",
      "  Rscript plot_orf_method_comparison.R .\n",
      "  Rscript plot_orf_method_comparison.R path/to/SAMPLE_td_singlebest\n",
      "  Rscript plot_orf_method_comparison.R path/to/orf_method_comparison"
    )
  }

  raw <- args[[1]]
  if (raw %in% c(".", "./")) {
    cand <- getwd()
    if (!has_td_layout(cand)) {
      stop("Current directory does not contain complete/ and with_partials/ summaries.")
    }
    return(normalizePath(cand, winslash = "/", mustWork = TRUE))
  }

  p <- normalizePath(raw, winslash = "/", mustWork = TRUE)
  if (has_td_layout(p)) {
    return(p)
  }

  bn <- basename(p)
  if (tolower(bn) == "orf_method_comparison" || grepl("orf_method_comparison$", p, ignore.case = TRUE)) {
    hits <- find_single_td_singlebest(p)
    if (length(hits) == 0L) {
      stop("No *_td_singlebest subfolder found under: ", p)
    }
    if (length(hits) > 1L) {
      stop(
        "Multiple *_td_singlebest folders under:\n  ", p,
        "\nPass the full path to one run, e.g.:\n  ",
        hits[[1]]
      )
    }
    message("Using: ", hits[[1]])
    return(hits[[1]])
  }

  stop(
    "Not a valid layout: ", p,
    "\nExpected complete/orf_comparison_summary.tsv and ",
    "with_partials/orf_comparison_summary.tsv\n",
    "Or pass the orf_method_comparison directory containing exactly one *_td_singlebest folder."
  )
}

args <- commandArgs(trailingOnly = TRUE)
base_dir <- resolve_base_dir(args)
sc_complete <- file.path(base_dir, "complete", "orf_comparison_summary.tsv")
sc_partials <- file.path(base_dir, "with_partials", "orf_comparison_summary.tsv")

if (!file.exists(sc_complete) || !file.exists(sc_partials)) {
  stop(
    "Expected:\n  ", sc_complete, "\n  ", sc_partials
  )
}

meuk_summary <- resolve_metaeuk_summary(base_dir, args)
if (!is.na(meuk_summary) && nzchar(meuk_summary) && file.exists(meuk_summary)) {
  message("MetaEuk summary: ", meuk_summary)
} else {
  meuk_summary <- NA_character_
  message("MetaEuk: no orf_comparison_summary.tsv (pass 2nd arg, ORF_SUMMARY_METAEUK, or sibling *_metaeuk/); plotting 3 methods only.")
}

td2_summary <- find_optional_summary(
  base_dir = base_dir,
  method_suffix = "td2",
  args = if (length(args) >= 3L) args[3] else character(0),
  env_var = "ORF_SUMMARY_TD2"
)
if (!is.na(td2_summary) && nzchar(td2_summary) && file.exists(td2_summary)) {
  message("TD2 summary: ", td2_summary)
} else {
  td2_summary <- NA_character_
  message("TD2: no orf_comparison_summary.tsv found (optional 3rd arg / ORF_SUMMARY_TD2 / sibling *_td2).")
}

kv_c <- read_summary_kv(sc_complete)
kv_p <- read_summary_kv(sc_partials)
kv_m <- if (!is.na(meuk_summary)) read_summary_kv(meuk_summary) else NULL
kv_td2 <- if (!is.na(td2_summary)) read_summary_kv(td2_summary) else NULL

# Six-frame metrics are duplicated in both summaries; read from complete run.
sf_in <- getn(kv_c, "diamond_input_queries_sixframe")
sf_hit_q <- getn(kv_c, "diamond_hit_queries_sixframe")
sf_n_c <- getn(kv_c, "n_contigs_sixframe")
sf_mean_len <- getn(kv_c, "sixframe_mean_aa_len")

tdc_in <- getn(kv_c, "diamond_input_queries_transdecoder")
tdc_hit_q <- getn(kv_c, "diamond_hit_queries_transdecoder")
tdc_n_c <- getn(kv_c, "n_contigs_transdecoder")
tdc_mean_len <- getn(kv_c, "td_longest_per_contig_mean_aa_len")

tdp_in <- getn(kv_p, "diamond_input_queries_transdecoder")
tdp_hit_q <- getn(kv_p, "diamond_hit_queries_transdecoder")
tdp_n_c <- getn(kv_p, "n_contigs_transdecoder")
tdp_mean_len <- getn(kv_p, "td_longest_per_contig_mean_aa_len")

muk_in <- muk_hit_q <- muk_n_c <- muk_mean_len <- NA_real_
if (!is.null(kv_m)) {
  muk_in <- getn(kv_m, "diamond_input_queries_transdecoder")
  muk_hit_q <- getn(kv_m, "diamond_hit_queries_transdecoder")
  muk_n_c <- getn(kv_m, "n_contigs_transdecoder")
  muk_mean_len <- getn(kv_m, "td_longest_per_contig_mean_aa_len")
}

td2_in <- td2_hit_q <- td2_n_c <- td2_mean_len <- NA_real_
if (!is.null(kv_td2)) {
  td2_in <- getn(kv_td2, "diamond_input_queries_transdecoder")
  td2_hit_q <- getn(kv_td2, "diamond_hit_queries_transdecoder")
  td2_n_c <- getn(kv_td2, "n_contigs_transdecoder")
  td2_mean_len <- getn(kv_td2, "td_longest_per_contig_mean_aa_len")
}

safe_rate <- function(num, den) {
  if (is.na(num) || is.na(den) || den <= 0) {
    return(NA_real_)
  }
  100 * num / den
}

plot_df <- bind_rows(
  tibble(
    method = "Six-frame\n(best ORF / contig)",
    marferret_hit_pct = safe_rate(sf_hit_q, sf_in),
    td_contigs_pct_of_sf = 100,
    mean_peptide_aa = sf_mean_len,
    n_input_orfs = sf_in,
    n_contigs = sf_n_c
  ),
  tibble(
    method = "TransDecoder\nsingle-best, complete proteins only",
    marferret_hit_pct = safe_rate(tdc_hit_q, tdc_in),
    td_contigs_pct_of_sf = safe_rate(tdc_n_c, sf_n_c),
    mean_peptide_aa = tdc_mean_len,
    n_input_orfs = tdc_in,
    n_contigs = tdc_n_c
  ),
  tibble(
    method = "TransDecoder\nsingle-best, complete+incomplete proteins",
    marferret_hit_pct = safe_rate(tdp_hit_q, tdp_in),
    td_contigs_pct_of_sf = safe_rate(tdp_n_c, sf_n_c),
    mean_peptide_aa = tdp_mean_len,
    n_input_orfs = tdp_in,
    n_contigs = tdp_n_c
  )
)
if (!is.null(kv_m)) {
  plot_df <- bind_rows(
    plot_df,
    tibble(
      method = "MetaEuk\n(TARA easy-predict)",
      marferret_hit_pct = safe_rate(muk_hit_q, muk_in),
      td_contigs_pct_of_sf = safe_rate(muk_n_c, sf_n_c),
      mean_peptide_aa = muk_mean_len,
      n_input_orfs = muk_in,
      n_contigs = muk_n_c
    )
  )
}
if (!is.null(kv_td2)) {
  plot_df <- bind_rows(
    plot_df,
    tibble(
      method = "TD2\n(single-best proteins)",
      marferret_hit_pct = safe_rate(td2_hit_q, td2_in),
      td_contigs_pct_of_sf = safe_rate(td2_n_c, sf_n_c),
      mean_peptide_aa = td2_mean_len,
      n_input_orfs = td2_in,
      n_contigs = td2_n_c
    )
  )
}

method_levels <- c(
  "Six-frame\n(best ORF / contig)",
  "TransDecoder\nsingle-best, complete proteins only",
  "TransDecoder\nsingle-best, complete+incomplete proteins"
)
if (!is.null(kv_m)) {
  method_levels <- c(method_levels, "MetaEuk\n(TARA easy-predict)")
}
if (!is.null(kv_td2)) {
  method_levels <- c(method_levels, "TD2\n(single-best proteins)")
}

method_levels_used <- method_levels[method_levels %in% unique(plot_df$method)]
plot_df <- plot_df |>
  mutate(method = factor(.data$method, levels = method_levels_used))

long_df <- plot_df |>
  pivot_longer(
    cols = c(marferret_hit_pct, td_contigs_pct_of_sf, mean_peptide_aa, n_input_orfs, n_contigs),
    names_to = "metric",
    values_to = "value"
  ) |>
  mutate(
    metric_label = factor(
      metric,
      levels = c(
        "marferret_hit_pct",
        "td_contigs_pct_of_sf",
        "mean_peptide_aa",
        "n_input_orfs",
        "n_contigs"
      ),
      labels = c(
        "MarFERReT hit rate (% of Diamond queries)",
        "Caller contigs as % of six-frame contigs",
        "Mean peptide length (aa)",
        "Diamond input sequences (queries)",
        "Contigs (six-frame) or caller parent contigs"
      )
    )
  )

fill_colors <- c(
  "Six-frame\n(best ORF / contig)" = "#4A6FA5",
  "TransDecoder\nsingle-best, complete proteins only" = "#C25552",
  "TransDecoder\nsingle-best, complete+incomplete proteins" = "#2A9D8F",
  "MetaEuk\n(TARA easy-predict)" = "#8E5AA8",
  "TD2\n(single-best proteins)" = "#D98E04"
)
fill_colors <- fill_colors[levels(plot_df$method)]

p <- ggplot(long_df, aes(x = method, y = value, fill = method)) +
  geom_col(width = 0.78, show.legend = FALSE) +
  facet_wrap(~metric_label, scales = "free_y", ncol = 2) +
  scale_fill_manual(values = fill_colors, drop = FALSE) +
  labs(
    title = "ORF prediction approaches vs MarFERReT (Diamond blastp)",
    subtitle = paste0(
      "Higher hit % = more queries with a MarFERReT hit at the run e-value; ",
      "not biological accuracy. MetaEuk / TransDecoder vs same six-frame assembly."
    ),
    x = NULL,
    y = NULL
  ) +
  theme_minimal(base_size = 11) +
  theme(
    axis.text.x = element_text(
      angle = if (!is.null(kv_m)) 35 else 0,
      hjust = if (!is.null(kv_m)) 1 else 0.5,
      size = 9
    ),
    strip.text = element_text(size = 9, face = "plain"),
    plot.title = element_text(face = "bold")
  )

out_prefix <- Sys.getenv("ORF_PLOT_OUT", unset = "")
if (!nzchar(out_prefix)) {
  out_prefix <- file.path(base_dir, "orf_method_comparison_plot")
}

plot_w <- if (!is.null(kv_m)) 12 else 10
ggsave(paste0(out_prefix, ".pdf"), p, width = plot_w, height = 8.5)
ggsave(paste0(out_prefix, ".png"), p, width = plot_w, height = 8.5, dpi = 150, bg = "white")

message("Wrote:\n  ", out_prefix, ".pdf\n  ", out_prefix, ".png")

# Compact numeric table for the report / supplementary
tbl_out <- file.path(base_dir, "orf_method_comparison_plot_table.tsv")
plot_df |>
  mutate(
    marferret_hit_pct = round(marferret_hit_pct, 2),
    td_contigs_pct_of_sf = round(td_contigs_pct_of_sf, 2),
    mean_peptide_aa = round(mean_peptide_aa, 2)
  ) |>
  write_tsv(tbl_out)
message("Wrote: ", tbl_out)