#!/usr/bin/env Rscript
# ============================================================================
# BioF3 Module 02: R and ggplot2 with Real PBMC 3k Data
# ============================================================================
#
# This script uses the public 10x Genomics PBMC 3k filtered matrix.
# It generates ggplot2 teaching figures from real counts and QC metrics.
#
# Usage:
#   Rscript scripts/module02_complete_sci.R
#
# Optional environment variables:
#   BIOF3_DATA_DIR=/path/to/data
#   BIOF3_OUTPUT_DIR=/path/to/static/img/tutorial/single-cell/module02
#
# ============================================================================

options(stringsAsFactors = FALSE)

args <- commandArgs(trailingOnly = FALSE)
file_arg <- grep("^--file=", args, value = TRUE)
script_path <- if (length(file_arg) > 0) {
  normalizePath(sub("^--file=", "", file_arg[[1]]), mustWork = TRUE)
} else {
  normalizePath("scripts/single-cell/sc02_cellranger_sci.R", mustWork = FALSE)
}

script_dir <- dirname(script_path)
project_root <- if (basename(script_dir) == "scripts" && basename(dirname(script_dir)) == "static") {
  dirname(dirname(script_dir))
} else {
  dirname(script_dir)
}

data_root <- Sys.getenv("BIOF3_DATA_DIR", file.path(path.expand("~"), "biof3-data"))
pbmc_dir <- file.path(data_root, "pbmc3k")
output_dir <- Sys.getenv(
  "BIOF3_OUTPUT_DIR",
  file.path(project_root, "static", "img", "tutorial", "modules", "module02")
)

dir.create(pbmc_dir, recursive = TRUE, showWarnings = FALSE)
dir.create(output_dir, recursive = TRUE, showWarnings = FALSE)

pbmc_url <- paste0(
  "https://cf.10xgenomics.com/samples/cell-exp/1.1.0/pbmc3k/",
  "pbmc3k_filtered_gene_bc_matrices.tar.gz"
)
pbmc_tar <- file.path(pbmc_dir, "pbmc3k_filtered_gene_bc_matrices.tar.gz")
matrix_dir <- file.path(pbmc_dir, "filtered_gene_bc_matrices", "hg19")
matrix_file <- file.path(matrix_dir, "matrix.mtx")
genes_file <- file.path(matrix_dir, "genes.tsv")
barcodes_file <- file.path(matrix_dir, "barcodes.tsv")

required_packages <- c("Matrix", "ggplot2", "dplyr", "tidyr", "patchwork", "scales", "pheatmap")
for (pkg in required_packages) {
  if (!requireNamespace(pkg, quietly = TRUE)) {
    install.packages(pkg, repos = "https://cloud.r-project.org/")
  }
}

library(Matrix)
library(ggplot2)
library(dplyr)
library(tidyr)
library(patchwork)
library(scales)
library(pheatmap)

biof3_colors <- c(
  green = "#0f766e",
  mint = "#43d1ae",
  blue = "#2563eb",
  amber = "#f59e0b",
  red = "#dc2626",
  purple = "#7c3aed",
  slate = "#475569",
  pale = "#e8f5ef"
)

theme_biof3 <- function(base_size = 12) {
  theme_classic(base_size = base_size) +
    theme(
      axis.text = element_text(color = "#111827"),
      axis.title = element_text(color = "#111827", face = "bold"),
      plot.title = element_text(color = "#12342f", face = "bold", hjust = 0),
      plot.subtitle = element_text(color = "#526760", hjust = 0),
      legend.title = element_text(face = "bold"),
      legend.position = "right",
      strip.background = element_blank(),
      strip.text = element_text(color = "#12342f", face = "bold")
    )
}

download_pbmc3k <- function() {
  if (!file.exists(pbmc_tar)) {
    message("Downloading PBMC 3k data from 10x Genomics...")
    download.file(pbmc_url, destfile = pbmc_tar, mode = "wb", quiet = FALSE)
  }

  if (!file.exists(matrix_file)) {
    message("Extracting PBMC 3k archive...")
    utils::untar(pbmc_tar, exdir = pbmc_dir)
  }
}

read_pbmc3k <- function() {
  download_pbmc3k()

  counts <- Matrix::readMM(matrix_file)
  genes <- read.delim(genes_file, header = FALSE, sep = "\t")
  barcodes <- read.delim(barcodes_file, header = FALSE, sep = "\t")

  rownames(counts) <- make.unique(genes[[2]])
  colnames(counts) <- barcodes[[1]]
  as(counts, "CsparseMatrix")
}

message("=== Loading real PBMC 3k data ===")
counts <- read_pbmc3k()
message("Genes: ", nrow(counts))
message("Cells: ", ncol(counts))

mt_genes <- grepl("^MT-", rownames(counts))
n_counts <- Matrix::colSums(counts)
n_features <- Matrix::colSums(counts > 0)
percent_mt <- Matrix::colSums(counts[mt_genes, , drop = FALSE]) / n_counts * 100

qc <- data.frame(
  cell = colnames(counts),
  nCount_RNA = as.numeric(n_counts),
  nFeature_RNA = as.numeric(n_features),
  percent_mt = as.numeric(percent_mt)
)

gene_summary <- data.frame(
  gene = rownames(counts),
  total_counts = as.numeric(Matrix::rowSums(counts)),
  detected_cells = as.numeric(Matrix::rowSums(counts > 0))
) %>%
  mutate(mean_counts = total_counts / ncol(counts))

marker_genes <- c("IL7R", "CCR7", "S100A8", "S100A9", "MS4A1", "CD79A", "NKG7", "GNLY", "PPBP")
marker_genes <- marker_genes[marker_genes %in% rownames(counts)]

gene_expression <- function(gene) {
  data.frame(
    cell = colnames(counts),
    gene = gene,
    counts = as.numeric(counts[gene, ]),
    log_counts = log1p(as.numeric(counts[gene, ]))
  )
}

marker_long <- bind_rows(lapply(marker_genes, gene_expression))

# ============================================================================
# Figure 1: Bar plot - real top expressed genes
# ============================================================================

top_genes <- gene_summary %>%
  filter(!grepl("^MT-|^RPL|^RPS", gene)) %>%
  arrange(desc(total_counts)) %>%
  slice_head(n = 10) %>%
  mutate(gene = factor(gene, levels = rev(gene)))

p1 <- ggplot(top_genes, aes(x = gene, y = total_counts)) +
  geom_col(fill = biof3_colors["green"], width = 0.72) +
  coord_flip() +
  scale_y_continuous(labels = comma) +
  labs(
    title = "Top expressed genes in PBMC 3k",
    subtitle = "Bar plot built from the real 10x filtered matrix",
    x = NULL,
    y = "Total UMI counts"
  ) +
  theme_biof3()

ggsave(file.path(output_dir, "01-bar-plot.png"), p1, width = 7.5, height = 5, dpi = 300, bg = "white")

# ============================================================================
# Figure 2: Scatter plot - real QC relationship
# ============================================================================

p2 <- ggplot(qc, aes(x = nCount_RNA, y = nFeature_RNA, color = percent_mt)) +
  geom_point(alpha = 0.72, size = 1.2) +
  scale_x_continuous(labels = comma) +
  scale_y_continuous(labels = comma) +
  scale_color_gradient(low = biof3_colors["green"], high = biof3_colors["red"]) +
  labs(
    title = "Counts and detected genes per PBMC 3k cell",
    subtitle = "Each point is one real cell barcode",
    x = "Total UMI counts",
    y = "Detected genes",
    color = "Mitochondrial %"
  ) +
  theme_biof3()

ggsave(file.path(output_dir, "02-scatter-plot.png"), p2, width = 7, height = 5, dpi = 300, bg = "white")

# ============================================================================
# Figure 3: Box plot - real marker genes
# ============================================================================

selected_box_genes <- marker_genes[1:min(4, length(marker_genes))]
box_data <- marker_long %>%
  filter(gene %in% selected_box_genes)

p3 <- ggplot(box_data, aes(x = gene, y = log_counts, fill = gene)) +
  geom_boxplot(width = 0.6, outlier.shape = NA, alpha = 0.82) +
  scale_fill_manual(
    values = setNames(
      unname(rep(c(biof3_colors["green"], biof3_colors["blue"], biof3_colors["amber"], biof3_colors["purple"]), length.out = length(selected_box_genes))),
      selected_box_genes
    )
  ) +
  labs(
    title = "PBMC marker gene expression",
    subtitle = "Box plots show log1p UMI counts across real PBMC 3k cells",
    x = NULL,
    y = "log1p(UMI counts)"
  ) +
  theme_biof3() +
  theme(legend.position = "none")

ggsave(file.path(output_dir, "03-box-plot.png"), p3, width = 7, height = 5, dpi = 300, bg = "white")

# ============================================================================
# Figure 4: Histogram - real total counts
# ============================================================================

p4 <- ggplot(qc, aes(x = nCount_RNA)) +
  geom_histogram(bins = 55, fill = biof3_colors["blue"], color = "white", alpha = 0.84) +
  geom_vline(xintercept = median(qc$nCount_RNA), linetype = "dashed", color = biof3_colors["amber"], linewidth = 0.8) +
  scale_x_continuous(labels = comma) +
  labs(
    title = "PBMC 3k total UMI count distribution",
    subtitle = paste0("Median total counts = ", comma(round(median(qc$nCount_RNA)))),
    x = "Total UMI counts per cell",
    y = "Number of cells"
  ) +
  theme_biof3()

ggsave(file.path(output_dir, "04-histogram.png"), p4, width = 7, height = 5, dpi = 300, bg = "white")

# ============================================================================
# Figure 5: Violin plot - real marker genes
# ============================================================================

p5 <- ggplot(box_data, aes(x = gene, y = log_counts, fill = gene)) +
  geom_violin(trim = FALSE, alpha = 0.72) +
  geom_boxplot(width = 0.12, fill = "white", outlier.shape = NA) +
  scale_fill_manual(
    values = setNames(
      unname(rep(c(biof3_colors["green"], biof3_colors["blue"], biof3_colors["amber"], biof3_colors["purple"]), length.out = length(selected_box_genes))),
      selected_box_genes
    )
  ) +
  labs(
    title = "Distribution of marker gene expression",
    subtitle = "Violin plots reveal sparsity and cell-to-cell heterogeneity",
    x = NULL,
    y = "log1p(UMI counts)"
  ) +
  theme_biof3() +
  theme(legend.position = "none")

ggsave(file.path(output_dir, "05-violin-plot.png"), p5, width = 7, height = 5, dpi = 300, bg = "white")

# ============================================================================
# Figure 6: Theme comparison using real PBMC 3k data
# ============================================================================

base_plot <- ggplot(top_genes, aes(x = gene, y = total_counts)) +
  geom_col(fill = biof3_colors["green"], width = 0.72) +
  coord_flip() +
  scale_y_continuous(labels = comma) +
  labs(x = NULL, y = "Total UMI counts")

p6 <- (
  base_plot + labs(title = "Classic") + theme_classic(base_size = 11)
) | (
  base_plot + labs(title = "Minimal") + theme_minimal(base_size = 11)
) / (
  base_plot + labs(title = "BioF3") + theme_biof3(11)
) | (
  base_plot + labs(title = "Black & white") + theme_bw(base_size = 11)
)

ggsave(file.path(output_dir, "06-themes-comparison.png"), p6, width = 12, height = 8, dpi = 300, bg = "white")

# ============================================================================
# Figure 7: Facet plot - real marker genes
# ============================================================================

facet_genes <- marker_genes[1:min(6, length(marker_genes))]
facet_data <- marker_long %>%
  filter(gene %in% facet_genes)

p7 <- ggplot(facet_data, aes(x = log_counts)) +
  geom_histogram(bins = 35, fill = biof3_colors["green"], color = "white", alpha = 0.82) +
  facet_wrap(~ gene, scales = "free_y", ncol = 3) +
  labs(
    title = "Marker expression distributions in PBMC 3k",
    subtitle = "Facets compare real gene-level distributions across cells",
    x = "log1p(UMI counts)",
    y = "Cells"
  ) +
  theme_biof3()

ggsave(file.path(output_dir, "07-facet-plot.png"), p7, width = 10, height = 6, dpi = 300, bg = "white")

# ============================================================================
# Figure 8: Real expression distribution
# ============================================================================

p8 <- ggplot(marker_long, aes(x = gene, y = log_counts, fill = gene)) +
  geom_violin(trim = FALSE, alpha = 0.72) +
  geom_boxplot(width = 0.1, fill = "white", outlier.shape = NA) +
  labs(
    title = "PBMC 3k marker expression distribution",
    subtitle = "Real marker genes from the public 10x matrix",
    x = NULL,
    y = "log1p(UMI counts)"
  ) +
  theme_biof3() +
  theme(
    legend.position = "none",
    axis.text.x = element_text(angle = 30, hjust = 1)
  )

ggsave(file.path(output_dir, "08-expression-distribution.png"), p8, width = 9, height = 5.5, dpi = 300, bg = "white")

# ============================================================================
# Figure 9: Gene detection versus abundance
# ============================================================================

p9_data <- gene_summary %>%
  filter(detected_cells > 0) %>%
  mutate(marker = if_else(gene %in% marker_genes, "Marker gene", "Other gene"))

p9 <- ggplot(p9_data, aes(x = detected_cells, y = total_counts, color = marker)) +
  geom_point(alpha = 0.45, size = 1) +
  geom_point(
    data = filter(p9_data, marker == "Marker gene"),
    size = 2.5,
    alpha = 0.95
  ) +
  scale_x_log10(labels = comma) +
  scale_y_log10(labels = comma) +
  scale_color_manual(values = c("Marker gene" = biof3_colors["red"], "Other gene" = "grey70")) +
  labs(
    title = "Gene detection and abundance in PBMC 3k",
    subtitle = "A real public-data view of gene prevalence and total abundance",
    x = "Cells with detected expression (log10)",
    y = "Total UMI counts (log10)",
    color = NULL
  ) +
  theme_biof3()

ggsave(file.path(output_dir, "09-volcano-plot.png"), p9, width = 7, height = 5.5, dpi = 300, bg = "white")

# ============================================================================
# Figure 10: Heatmap from real PBMC 3k marker genes
# ============================================================================

heat_genes <- marker_genes[1:min(9, length(marker_genes))]
cell_order <- order(qc$nCount_RNA)
selected_cells <- colnames(counts)[cell_order[round(seq(1, length(cell_order), length.out = 90))]]
cell_totals <- qc$nCount_RNA[match(selected_cells, qc$cell)]

heat_counts <- as.matrix(counts[heat_genes, selected_cells, drop = FALSE])
heat_norm <- t(t(heat_counts) / cell_totals * 10000)
heat_log <- log1p(heat_norm)

png(file.path(output_dir, "10-heatmap.png"), width = 7, height = 6.5, units = "in", res = 300)
pheatmap(
  heat_log,
  color = colorRampPalette(c("#f8fafc", "#43d1ae", "#0f172a"))(100),
  cluster_rows = TRUE,
  cluster_cols = TRUE,
  show_colnames = FALSE,
  fontsize = 9,
  main = "PBMC 3k marker expression heatmap"
)
dev.off()

message("=== Generated Module 02 real-data figures ===")
print(list.files(output_dir, pattern = "\\.png$", full.names = FALSE))
message("Output directory: ", normalizePath(output_dir))
message("Data directory: ", normalizePath(pbmc_dir))
message("Done.")

sessionInfo()