#!/usr/bin/env Rscript
# ==============================================================================
# Chapter 10: R for Next Gen Biologist
# Script: r10_complexheatmaps_clustering.R
# Description: Advanced Omics Heatmaps & Clustering with ComplexHeatmap:
#              Row-wise Z-score normalization, hierarchical clustering algorithms,
#              multi-track clinical column annotations, and high-res export.
# Usage: Rscript r10_complexheatmaps_clustering.R
# ==============================================================================

cat("\n=======================================================\n")
cat(" OMICS.CO.IN - R FOR NEXT GEN BIOLOGIST: CHAPTER 10\n")
cat(" Clustered Omics Heatmaps & Multi-Track Annotations\n")
cat("=======================================================\n\n")

temp_dir <- tempdir()
has_ch <- requireNamespace("ComplexHeatmap", quietly = TRUE)
has_pheatmap <- requireNamespace("pheatmap", quietly = TRUE)

# 1. Simulating Co-regulated Gene Expression Matrix
set.seed(42)
n_genes <- 25
n_samples <- 12

genes <- c("TP53", "EGFR", "CDKN2A", "MYC", "KRAS", "PTEN", "BRCA1", "BRCA2", "PIK3CA", "ATM",
           paste0("TARGET_", 11:n_genes))
sample_ids <- paste0("Patient_", sprintf("%02d", 1:n_samples))

# Build modular cluster blocks
mat <- matrix(rnorm(n_genes * n_samples, mean = 5, sd = 1), nrow = n_genes, ncol = n_samples,
              dimnames = list(genes, sample_ids))

# Inject strong upregulated module in first 6 samples (e.g. Tumor Responders)
mat[1:5, 1:6] <- mat[1:5, 1:6] + 3.5
# Inject downregulated module in remaining 6 samples
mat[6:10, 7:12] <- mat[6:10, 7:12] - 3.2

# 2. Row Z-score Normalization (Z = (x - mean) / sd)
cat("[Step 1] Computing Row-wise Z-Scores...\n")
mat_z <- t(scale(t(mat)))
cat(sprintf("Matrix Standardized: Mean = %.2f | SD = %.2f\n", mean(mat_z), sd(mat_z)))

# 3. Clinical & Experimental Annotations
cat("\n[Step 2] Defining Multi-Track Sample Metadata...\n")
sample_anno <- data.frame(
  Subtype   = factor(rep(c("Basal-like", "Luminal A", "HER2+"), each = 4)),
  Response  = factor(rep(c("Responder", "Non-Responder"), each = 6)),
  Stage     = factor(rep(c("Stage_I", "Stage_II", "Stage_III", "Stage_IV"), times = 3)),
  Age       = round(runif(n_samples, 45, 75)),
  stringsAsFactors = FALSE
)
rownames(sample_anno) <- sample_ids
print(head(sample_anno, 6))

# 4. Rendering Heatmap
cat("\n[Step 3] Rendering Clustered Heatmap...\n")
output_png <- file.path(temp_dir, "complex_heatmap_genomics.png")

if (has_ch) {
  library(ComplexHeatmap, quietly = TRUE)
  library(circlize, quietly = TRUE)
  
  col_fun <- colorRamp2(c(-2, 0, 2), c("#2563EB", "#FFFFFF", "#DC2626"))
  
  ha_top <- HeatmapAnnotation(
    df = sample_anno,
    col = list(
      Subtype  = c("Basal-like" = "#E11D48", "Luminal A" = "#059669", "HER2+" = "#7C3AED"),
      Response = c("Responder" = "#10B981", "Non-Responder" = "#64748B")
    )
  )
  
  png(output_png, width = 2400, height = 2100, res = 300)
  ht <- Heatmap(
    mat_z,
    name = "Z-Score",
    col = col_fun,
    top_annotation = ha_top,
    clustering_distance_rows = "pearson",
    clustering_method_rows = "ward.D2",
    row_names_gp = gpar(fontsize = 9, fontface = "bold"),
    column_names_gp = gpar(fontsize = 8),
    column_title = "Oncogenic Transcriptomic Profile across Cohort"
  )
  draw(ht)
  dev.off()
  cat(sprintf("ComplexHeatmap output saved: %s\n", output_png))
  
} else if (has_pheatmap) {
  library(pheatmap, quietly = TRUE)
  png(output_png, width = 2400, height = 2100, res = 300)
  pheatmap(mat_z, annotation_col = sample_anno[, c("Subtype", "Response")],
           clustering_method = "ward.D2", main = "Multi-Track Heatmap (pheatmap)",
           color = colorRampPalette(c("#2563EB", "#FFFFFF", "#DC2626"))(100))
  dev.off()
  cat(sprintf("pheatmap output saved: %s\n", output_png))
  
} else {
  # Base R heatmap
  png(output_png, width = 2400, height = 2100, res = 300)
  heatmap(mat_z, scale = "none",
          col = colorRampPalette(c("#2563EB", "#FFFFFF", "#DC2626"))(100),
          main = "Clustered Heatmap (Base R)", margins = c(8, 8))
  dev.off()
  cat(sprintf("Base R heatmap output saved: %s\n", output_png))
}

cat("\n[SUCCESS] Chapter 10 heatmap clustering script executed successfully.\n")