#!/usr/bin/env Rscript
# =============================================================================
# make_figures.R
# Publication-style viromics figures in R (ggplot2 / pheatmap / vegan).
# Reads the teaching tables in data/example/ so it runs without the full
# pipeline. Point the paths at your own Phase 4 outputs for real figures.
#
# Install once:
#   install.packages(c("ggplot2","pheatmap","vegan","readr","dplyr","tidyr"))
#
# Usage:  Rscript make_figures.R [INPUT_DIR] [OUTPUT_DIR]
# =============================================================================
suppressMessages({
  library(ggplot2); library(readr); library(dplyr); library(tidyr)
})

args <- commandArgs(trailingOnly = TRUE)
IN  <- ifelse(length(args) >= 1, args[1], "data/example")
OUT <- ifelse(length(args) >= 2, args[2], "visualization/figures")
dir.create(OUT, showWarnings = FALSE, recursive = TRUE)

# Codanics house palette
teal <- "#008b8b"; navy <- "#05043b"; terra <- "#c0432a"; gold <- "#d9a521"
quality_cols <- c("Complete"=navy, "High-quality"=teal,
                  "Medium-quality"=gold, "Low-quality"="#8a97a0")

theme_codanics <- theme_minimal(base_size = 13) +
  theme(plot.title = element_text(face = "bold", colour = navy),
        panel.grid.minor = element_blank())

# ---- Figure 1: CheckV quality ----------------------------------------------
checkv <- read_tsv(file.path(IN, "checkv_quality_summary.tsv"),
                   show_col_types = FALSE)
checkv$checkv_quality <- factor(
  checkv$checkv_quality,
  levels = c("Complete","High-quality","Medium-quality","Low-quality"))
g1 <- ggplot(checkv, aes(checkv_quality, fill = checkv_quality)) +
  geom_bar() +
  scale_fill_manual(values = quality_cols, guide = "none") +
  labs(title = "CheckV quality of recovered viral genomes",
       x = NULL, y = "Number of vOTUs") + theme_codanics
ggsave(file.path(OUT, "fig-checkv-quality.png"), g1, width = 7, height = 4.2, dpi = 300)

# ---- Figure 2: mean abundance per vOTU -------------------------------------
ab <- read_tsv(file.path(IN, "votu_abundance.tsv"), show_col_types = FALSE)
ab_long <- pivot_longer(ab, -votu, names_to = "sample", values_to = "tpm")
mean_ab <- ab_long %>% group_by(votu) %>% summarise(mean_tpm = mean(tpm))
g2 <- ggplot(mean_ab, aes(reorder(votu, mean_tpm), mean_tpm)) +
  geom_col(fill = navy) + coord_flip() +
  labs(title = "Mean abundance per vOTU (TPM)", x = NULL, y = "Mean TPM") +
  theme_codanics
ggsave(file.path(OUT, "fig-top-abundance.png"), g2, width = 7, height = 4.2, dpi = 300)

# ---- Figure 3: taxonomy by family ------------------------------------------
tax <- read_tsv(file.path(IN, "taxonomy.tsv"), show_col_types = FALSE)
g3 <- ggplot(tax, aes(forcats::fct_infreq(family))) +
  geom_bar(fill = teal) +
  labs(title = "vOTU taxonomy by family (geNomad)", x = NULL, y = "Number of vOTUs") +
  theme_codanics + theme(axis.text.x = element_text(angle = 30, hjust = 1))
ggsave(file.path(OUT, "fig-taxonomy.png"), g3, width = 7, height = 4.2, dpi = 300)

# ---- Figure 4: abundance heatmap (pheatmap) --------------------------------
if (requireNamespace("pheatmap", quietly = TRUE)) {
  mat <- as.matrix(ab[,-1]); rownames(mat) <- ab$votu
  meta <- read_tsv(file.path(IN, "sample_metadata.tsv"), show_col_types = FALSE)
  ann <- data.frame(treatment = meta$treatment, row.names = meta$sample)
  pheatmap::pheatmap(log10(mat + 1), annotation_col = ann,
                     color = colorRampPalette(c("#f7fbfb", teal, navy))(50),
                     main = "vOTU abundance (log10 TPM+1)",
                     filename = file.path(OUT, "fig-abundance-heatmap.png"),
                     width = 6.5, height = 5)
}

# ---- Figure 5: alpha diversity (vegan) -------------------------------------
if (requireNamespace("vegan", quietly = TRUE)) {
  comm <- t(as.matrix(ab[,-1])); colnames(comm) <- ab$votu
  richness <- rowSums(comm > 0)
  shannon  <- vegan::diversity(comm, index = "shannon")
  div <- data.frame(sample = names(richness), richness, shannon)
  g5 <- ggplot(div, aes(sample, richness)) + geom_col(fill = teal) +
    labs(title = "Observed vOTU richness per sample", x = NULL, y = "Observed vOTUs") +
    theme_codanics + theme(axis.text.x = element_text(angle = 20, hjust = 1))
  ggsave(file.path(OUT, "fig-alpha-diversity.png"), g5, width = 6.5, height = 4.2, dpi = 300)
}

cat("Figures written to", OUT, "\n")