#!/usr/bin/env Rscript
# ============================================================
# MuSiC cell type deconvolution
# Reference: Consensus-WMB-Macosko  (mouse whole brain)
#
# Run prep_music_ref.py first to generate the sc reference files.
#
# Usage:
#   Rscript run_music.R \
#     --sc_dir  /BRC/yan/BPA/sc_brain_ref/MuSiC_ref_subclass_maxpd10 \
#     --bulk    /path/to/bulk_counts.txt \
#     --outdir  /path/to/output
# ============================================================

suppressPackageStartupMessages({
  library(MuSiC)
  library(SingleCellExperiment)
  library(Matrix)
  library(Biobase)
  library(optparse)
})

# ── args ──────────────────────────────────────────────────────────────────────
opt_list <- list(
  make_option("--sc_dir",  type = "character",
              help = "Directory from prep_music_ref.py"),
  make_option("--bulk",    type = "character",
              help = "Bulk RNA-seq count matrix (genes x samples, tab-delimited, header + rownames)"),
  make_option("--outdir",  type = "character", default = NULL,
              help = "Output directory [default: same as --sc_dir]"),
  make_option("--celltype_col", type = "character", default = "celltype",
              help = "Column name for cell types in sc_meta.csv"),
  make_option("--donor_col",    type = "character", default = "donor",
              help = "Column name for donors in sc_meta.csv")
)
opt <- parse_args(OptionParser(option_list = opt_list))

if (is.null(opt$sc_dir) || is.null(opt$bulk)) {
  stop("--sc_dir and --bulk are required")
}
if (is.null(opt$outdir)) opt$outdir <- opt$sc_dir
dir.create(opt$outdir, showWarnings = FALSE, recursive = TRUE)

# ── 1. load single-cell reference ─────────────────────────────────────────────
cat("[1] Loading single-cell reference ...\n")

counts_mtx <- readMM(file.path(opt$sc_dir, "sc_counts.mtx.gz"))
gene_names  <- readLines(file.path(opt$sc_dir, "sc_genes.txt"))
sc_meta     <- read.csv(file.path(opt$sc_dir, "sc_meta.csv"),
                        row.names = 1, stringsAsFactors = FALSE)

rownames(counts_mtx) <- gene_names
colnames(counts_mtx) <- rownames(sc_meta)

cat("    Genes :", nrow(counts_mtx), "\n")
cat("    Cells :", ncol(counts_mtx), "\n")
cat("    Cell types :",
    length(unique(sc_meta[[opt$celltype_col]])), "\n")
cat("    Donors :",
    length(unique(sc_meta[[opt$donor_col]])), "\n")

# cell type × donor summary
cat("\n    Cell type distribution:\n")
ct_tab <- sort(table(sc_meta[[opt$celltype_col]]), decreasing = TRUE)
print(ct_tab)

# ── 2. build SingleCellExperiment ─────────────────────────────────────────────
cat("\n[2] Building SingleCellExperiment ...\n")

sce <- SingleCellExperiment(
  assays   = list(counts = counts_mtx),
  colData  = sc_meta
)

cat("    SCE built:", dim(sce)[1], "genes x", dim(sce)[2], "cells\n")

# save SCE for reuse
sce_path <- file.path(opt$outdir, "sc_sce.rds")
saveRDS(sce, sce_path)
cat("    Saved SCE to:", sce_path, "\n")

# ── 3. load bulk RNA-seq ───────────────────────────────────────────────────────
cat("\n[3] Loading bulk RNA-seq ...\n")

bulk_raw <- read.table(opt$bulk, header = TRUE, sep = "\t",
                       row.names = 1, check.names = FALSE)
bulk_mat <- as.matrix(bulk_raw)

cat("    Bulk matrix:", nrow(bulk_mat), "genes x", ncol(bulk_mat), "samples\n")
cat("    Samples:", colnames(bulk_mat), "\n")

# build ExpressionSet for MuSiC
bulk_eset <- ExpressionSet(assayData = bulk_mat)

# ── 4. find common genes ───────────────────────────────────────────────────────
common_genes <- intersect(rownames(sce), rownames(bulk_mat))
cat("\n[4] Common genes between SC ref and bulk:", length(common_genes), "\n")

if (length(common_genes) < 1000) {
  warning("Fewer than 1000 common genes — check gene symbol format")
}

# ── 5. run MuSiC ──────────────────────────────────────────────────────────────
cat("\n[5] Running MuSiC deconvolution ...\n")

music_result <- music_prop(
  bulk.mtx  = bulk_mat[common_genes, ],
  sc.sce    = sce[common_genes, ],
  clusters  = opt$celltype_col,
  samples   = opt$donor_col,
  select.ct = NULL,          # use all cell types
  verbose   = TRUE
)

# ── 6. save results ───────────────────────────────────────────────────────────
cat("\n[6] Saving results ...\n")

# estimated proportions (samples x cell types)
prop_df <- as.data.frame(music_result$Est.prop.weighted)
prop_df$sample <- rownames(prop_df)
prop_df <- prop_df[, c("sample", setdiff(colnames(prop_df), "sample"))]

write.csv(prop_df,
          file.path(opt$outdir, "MuSiC_proportions.csv"),
          row.names = FALSE)

# also save weight matrix
weight_df <- as.data.frame(music_result$Weight.gene)
write.csv(weight_df,
          file.path(opt$outdir, "MuSiC_gene_weights.csv"))

# ── 7. summary plot ───────────────────────────────────────────────────────────
cat("\n[7] Plotting proportions ...\n")

pdf(file.path(opt$outdir, "MuSiC_proportions.pdf"),
    width = max(8, ncol(prop_df) * 0.3), height = 6)

prop_mat <- as.matrix(prop_df[, -1])
rownames(prop_mat) <- prop_df$sample

barplot(t(prop_mat),
        col     = rainbow(ncol(prop_mat)),
        las     = 2,
        legend  = colnames(prop_mat),
        args.legend = list(x = "topright", cex = 0.6, bty = "n"),
        ylab    = "Estimated proportion",
        main    = "MuSiC cell type proportions",
        border  = NA)
dev.off()

# ── print summary ──────────────────────────────────────────────────────────────
cat("\n", strrep("=", 60), "\n")
cat("Done.\n")
cat("  Proportions  :", file.path(opt$outdir, "MuSiC_proportions.csv"), "\n")
cat("  Gene weights :", file.path(opt$outdir, "MuSiC_gene_weights.csv"), "\n")
cat("  Plot         :", file.path(opt$outdir, "MuSiC_proportions.pdf"), "\n")
cat("  SCE (reuse)  :", sce_path, "\n")
cat(strrep("=", 60), "\n")

cat("\nMean proportions across samples:\n")
print(round(colMeans(prop_mat), 4))