#!/usr/bin/env Rscript

suppressPackageStartupMessages({
  library(MuSiC)
  library(SingleCellExperiment)
  library(Biobase)
  library(S4Vectors)
  library(Matrix)
  library(dplyr)
  library(tidyr)
  library(tibble)
  library(ggplot2)
  library(ggrepel)
})

# ── 0. Results directory tree & logger ──────────────────────────────────────

RESULTS_DIR <- "/BRC/yan/BPA/sc_brain_ref/music_result_subclass/"

DIRS <- list(
  s01  = file.path(RESULTS_DIR, "step01_data_loading"),
  s02  = file.path(RESULTS_DIR, "step02_group_subsetting"),
  s03  = file.path(RESULTS_DIR, "step03_reference"),
  s04  = file.path(RESULTS_DIR, "step04_qc"),
  s05  = file.path(RESULTS_DIR, "step05_music"),
  s06  = file.path(RESULTS_DIR, "step06_metadata"),
  s07  = file.path(RESULTS_DIR, "step07_proportions"),
  s08  = file.path(RESULTS_DIR, "step08_stats"),
  s09  = file.path(RESULTS_DIR, "step09_plots")
)

for (d in DIRS) dir.create(d, recursive = TRUE, showWarnings = FALSE)

LOG_FILE <- file.path(RESULTS_DIR, "run_log.txt")

cat(sprintf("=== Run started: %s ===\n\n",
            format(Sys.time(), "%Y-%m-%d %H:%M:%S")),
    file = LOG_FILE)

log_msg <- function(..., type = "INFO") {
  msg <- sprintf(
    "[%s] [%s] %s",
    format(Sys.time(), "%Y-%m-%d %H:%M:%S"),
    type,
    paste0(...)
  )
  message(msg)
  cat(msg, "\n", file = LOG_FILE, append = TRUE)
}

log_save <- function(path) {
  log_msg("Saved: ", path)
}

# ── 1. Load bulk count matrix ────────────────────────────────────────────────

log_msg("STEP 1 — Loading bulk count matrix")

bulk_file <- "/BRC/shuhua/target/BPA_UPenn/raw_data/RNA/rc_brain_adt"

stopifnot("Bulk count file not found — check path" = file.exists(bulk_file))

bulk_counts <- read.table(
  bulk_file,
  header = TRUE,
  row.names = 1,
  check.names = FALSE
)

bulk_counts <- as.data.frame(bulk_counts)

# Force numeric counts
bulk_counts[] <- lapply(bulk_counts, function(x) as.numeric(as.character(x)))

if (anyNA(bulk_counts)) {
  log_msg("WARNING: NA introduced when converting bulk counts to numeric", type = "WARN")
}

write.table(
  bulk_counts,
  file = "/BRC/yan/BPA/sc_brain_ref/music_result/raw_bulk_counts_br.tsv",
  sep = "\t",
  quote = FALSE,
  row.names = TRUE,
  col.names = NA
)

log_msg(sprintf("Raw bulk matrix: %d genes × %d samples",
                nrow(bulk_counts), ncol(bulk_counts)))

# Keep only BA lab samples
lab_prefix <- sub("_RNA_.*", "", colnames(bulk_counts))

log_msg("Lab sample counts before BA filter:")
print(table(lab_prefix))

ba_cols <- grep("^BA_RNA_", colnames(bulk_counts), value = TRUE)

if (length(ba_cols) == 0) {
  stop("No BA_RNA samples found in bulk count matrix.")
}

bulk_counts <- bulk_counts[, ba_cols, drop = FALSE]

log_msg(sprintf("After BA-only filter: %d genes × %d samples",
                nrow(bulk_counts), ncol(bulk_counts)))

cat("\nSample distribution after filter:\n")
cat("BA Ctrl  Male:   ", sum(grepl("^BA_RNA_Ctrl.*Br_M",    colnames(bulk_counts))), "\n")
cat("BA Ctrl  Female: ", sum(grepl("^BA_RNA_Ctrl.*Br_F",    colnames(bulk_counts))), "\n")
cat("BA 10mg  Male:   ", sum(grepl("^BA_RNA_BPA10mg.*Br_M", colnames(bulk_counts))), "\n")
cat("BA 10mg  Female: ", sum(grepl("^BA_RNA_BPA10mg.*Br_F", colnames(bulk_counts))), "\n")
cat("BA 10ug  Male:   ", sum(grepl("^BA_RNA_BPA10ug.*Br_M", colnames(bulk_counts))), "\n")
cat("BA 10ug  Female: ", sum(grepl("^BA_RNA_BPA10ug.*Br_F", colnames(bulk_counts))), "\n")

out <- file.path(DIRS$s01, "bulk_matrix_dim.txt")

writeLines(
  c(
    sprintf("Genes:   %d", nrow(bulk_counts)),
    sprintf("Samples: %d", ncol(bulk_counts)),
    "",
    "Sample names:",
    colnames(bulk_counts)
  ),
  out
)

log_save(out)

# ── 2. Subset bulk matrix by treatment group and sex ─────────────────────────

log_msg("STEP 2 — Subsetting bulk matrix by treatment x sex")

get_matrix <- function(treatment, sex) {
  sex_code <- ifelse(sex == "F", "Br_F", "Br_M")
  
  treat_cols <- grep(
    paste0("^BA_RNA_", treatment, ".*", sex_code),
    colnames(bulk_counts),
    value = TRUE
  )
  
  ctrl_cols <- grep(
    paste0("^BA_RNA_Ctrl.*", sex_code),
    colnames(bulk_counts),
    value = TRUE
  )
  
  if (length(ctrl_cols) == 0) {
    log_msg("WARNING: no control samples found for sex = ", sex, type = "WARN")
  }
  
  if (length(treat_cols) == 0) {
    log_msg("WARNING: no treatment samples found for ", treatment, "_", sex, type = "WARN")
  }
  
  bulk_counts[, c(ctrl_cols, treat_cols), drop = FALSE]
}

mat_10mg_F <- get_matrix("BPA10mg", "F")
mat_10mg_M <- get_matrix("BPA10mg", "M")
mat_10ug_F <- get_matrix("BPA10ug", "F")
mat_10ug_M <- get_matrix("BPA10ug", "M")

group_summary <- data.frame(
  group = c("BPA10mg_F", "BPA10mg_M", "BPA10ug_F", "BPA10ug_M"),
  n_samples = c(
    ncol(mat_10mg_F),
    ncol(mat_10mg_M),
    ncol(mat_10ug_F),
    ncol(mat_10ug_M)
  ),
  n_ctrl = c(
    sum(grepl("Ctrl", colnames(mat_10mg_F))),
    sum(grepl("Ctrl", colnames(mat_10mg_M))),
    sum(grepl("Ctrl", colnames(mat_10ug_F))),
    sum(grepl("Ctrl", colnames(mat_10ug_M)))
  ),
  n_treat = c(
    sum(grepl("BPA10mg", colnames(mat_10mg_F))),
    sum(grepl("BPA10mg", colnames(mat_10mg_M))),
    sum(grepl("BPA10ug", colnames(mat_10ug_F))),
    sum(grepl("BPA10ug", colnames(mat_10ug_M)))
  )
)

print(group_summary)

out <- file.path(DIRS$s02, "group_sample_counts.csv")
write.csv(group_summary, out, row.names = FALSE)
log_save(out)

# ── 3. Load MuSiC single-cell reference ──────────────────────────────────────

log_msg("STEP 3 — Loading single-cell reference")

library(Matrix)

SC_DIR <- "/BRC/yan/BPA/sc_brain_ref/MuSiC_ref_subclass_maxpd10"

counts_mtx <- readMM(
  pipe(paste("zcat", shQuote(file.path(SC_DIR, "sc_counts.mtx.gz")),
             "| sed '1s/unsigned-integer/integer/'"))
)
gene_names  <- readLines(file.path(SC_DIR, "sc_genes.txt"))
sc_meta     <- read.csv(file.path(SC_DIR, "sc_meta.csv"),
                        row.names = 1, stringsAsFactors = FALSE)

stopifnot("sc_counts row count != sc_genes line count" =
            nrow(counts_mtx) == length(gene_names))
stopifnot("sc_counts col count != sc_meta row count" =
            ncol(counts_mtx) == nrow(sc_meta))

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

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

log_msg(sprintf("SCE: %d genes x %d cells | %d cell types | %d donors",
                nrow(sce), ncol(sce),
                length(unique(sc_meta$celltype)),
                length(unique(sc_meta$donor))))

saveRDS(sce, file.path(DIRS$s03, "sc_sce.rds"))
log_save(file.path(DIRS$s03, "sc_sce.rds"))

# ── 4. QC — common genes ─────────────────────────────────────────────────────

log_msg("STEP 4 — Finding common genes between bulk and reference")

common_genes <- intersect(rownames(sce), rownames(bulk_counts))

log_msg(sprintf("Reference genes : %d", nrow(sce)))
log_msg(sprintf("Bulk genes      : %d", nrow(bulk_counts)))
log_msg(sprintf("Common genes    : %d", length(common_genes)))

if (length(common_genes) < 1000)
  log_msg("WARNING: fewer than 1000 common genes — check gene symbol format", type = "WARN")

writeLines(common_genes, file.path(DIRS$s04, "common_genes.txt"))
log_save(file.path(DIRS$s04, "common_genes.txt"))

# ── 5. Run MuSiC on all samples at once ──────────────────────────────────────
#
# Running on the full matrix avoids duplicating Ctrl samples across the four
# sub-matrices that get_matrix() produces (Ctrl_F appears in both mat_10mg_F
# and mat_10ug_F, causing Ctrl n to be doubled in downstream stats).

log_msg("STEP 5 — Running MuSiC deconvolution")

g <- intersect(rownames(sce), rownames(bulk_counts))
log_msg(sprintf("Running on full bulk matrix: %d samples, %d common genes",
                ncol(bulk_counts), length(g)))

res_all <- music_prop(
  bulk.mtx  = as.matrix(bulk_counts[g, ]),
  sc.sce    = sce[g, ],
  clusters  = "celltype",
  samples   = "donor",
  select.ct = NULL,
  verbose   = FALSE
)

# Derive treatment and sex from sample names — avoids mislabelling Ctrl rows
props_all <- as.data.frame(res_all$Est.prop.weighted)
props_all$sample <- rownames(props_all)
props_all$treatment <- dplyr::case_when(
  grepl("BPA10mg", props_all$sample) ~ "BPA10mg",
  grepl("BPA10ug", props_all$sample) ~ "BPA10ug",
  grepl("Ctrl",    props_all$sample) ~ "Ctrl",
  TRUE ~ NA_character_
)
props_all$sex   <- ifelse(grepl("Br_F", props_all$sample), "F", "M")
props_all$group <- ifelse(props_all$treatment == "Ctrl",
                          paste0("Ctrl_",         props_all$sex),
                          paste0(props_all$treatment, "_", props_all$sex))

cell_cols_order <- setdiff(colnames(props_all),
                           c("sample", "group", "treatment", "sex"))
props_all <- props_all[, c("sample", "group", "treatment", "sex",
                            cell_cols_order)]

out <- file.path(DIRS$s05, "proportions_all_samples.csv")
write.csv(props_all, out, row.names = FALSE)
log_save(out)

log_msg(sprintf("MuSiC complete: %d samples x %d cell types",
                nrow(props_all), length(cell_cols_order)))
log_msg("Treatment breakdown:")
print(table(props_all$treatment, props_all$sex))

log_msg("=== MuSiC done — proceeding to summary steps ===")

# ── 6. Metadata enrichment ────────────────────────────────────────────────────

log_msg("STEP 6 — Enriching proportions with sample metadata")

cell_cols <- cell_cols_order   # defined in Step 5

# long format: one row per sample × cell type
props_long <- props_all %>%
  pivot_longer(cols = all_of(cell_cols),
               names_to  = "celltype",
               values_to = "proportion") %>%
  mutate(
    dose = case_when(
      treatment == "Ctrl"    ~ 0,
      treatment == "BPA10ug" ~ 10,
      treatment == "BPA10mg" ~ 10000
    )
  )

out <- file.path(DIRS$s06, "proportions_long.csv")
write.csv(props_long, out, row.names = FALSE)
log_save(out)

# ── 7. Per-group proportion summaries ────────────────────────────────────────

log_msg("STEP 7 — Summarising proportions per group")

prop_summary <- props_long %>%
  group_by(group, treatment, sex, celltype) %>%
  summarise(
    n        = n(),
    mean_prop = mean(proportion),
    sd_prop   = sd(proportion),
    median_prop = median(proportion),
    .groups = "drop"
  )

out <- file.path(DIRS$s07, "proportion_summary.csv")
write.csv(prop_summary, out, row.names = FALSE)
log_save(out)

# mean proportion table: cell types × groups (wide)
prop_wide <- prop_summary %>%
  select(group, celltype, mean_prop) %>%
  pivot_wider(names_from = group, values_from = mean_prop) %>%
  arrange(celltype)

out <- file.path(DIRS$s07, "proportion_mean_wide.csv")
write.csv(prop_wide, out, row.names = FALSE)
log_save(out)

# ── 8. Statistics: BPA vs Ctrl per cell type per sex ─────────────────────────

log_msg("STEP 8 — Wilcoxon tests: BPA vs Ctrl per cell type x sex")

run_wilcox <- function(df, sex_filter, treat_a, label) {
  df_sub <- df %>%
    filter(sex == sex_filter, treatment %in% c("Ctrl", treat_a))
  df_sub %>%
    group_by(celltype) %>%
    summarise(
      sex        = sex_filter,
      comparison = label,
      n_ctrl     = sum(treatment == "Ctrl"),
      n_treat    = sum(treatment == treat_a),
      mean_ctrl  = mean(proportion[treatment == "Ctrl"],  na.rm = TRUE),
      mean_treat = mean(proportion[treatment == treat_a], na.rm = TRUE),
      log2FC     = log2((mean_treat + 1e-6) / (mean_ctrl + 1e-6)),
      p_value    = tryCatch(
        wilcox.test(proportion[treatment == treat_a],
                    proportion[treatment == "Ctrl"])$p.value,
        error = function(e) NA_real_
      ),
      .groups = "drop"
    )
}

stat_results <- bind_rows(
  run_wilcox(props_long, "F", "BPA10ug", "BPA10ug_vs_Ctrl"),
  run_wilcox(props_long, "M", "BPA10ug", "BPA10ug_vs_Ctrl"),
  run_wilcox(props_long, "F", "BPA10mg", "BPA10mg_vs_Ctrl"),
  run_wilcox(props_long, "M", "BPA10mg", "BPA10mg_vs_Ctrl")
) %>%
  group_by(sex, comparison) %>%
  mutate(p_adj = p.adjust(p_value, method = "BH")) %>%
  ungroup() %>%
  arrange(p_adj)

out <- file.path(DIRS$s08, "wilcox_BPA_vs_Ctrl.csv")
write.csv(stat_results, out, row.names = FALSE)
log_save(out)

n_sig <- sum(stat_results$p_adj < 0.05, na.rm = TRUE)
log_msg(sprintf("Significant cell types (BH-adj p<0.05): %d", n_sig))

# ── 9. Plots ─────────────────────────────────────────────────────────────────

log_msg("STEP 9 — Generating plots")

# colour palette for cell types
n_ct   <- length(cell_cols)
ct_pal <- setNames(
  colorRampPalette(c("#4DBBD5","#E64B35","#00A087","#3C5488",
                     "#F39B7F","#8491B4","#91D1C2","#DC0000"))(n_ct),
  cell_cols
)

# 9a — stacked bar: mean proportion per group
p_bar <- prop_summary %>%
  ggplot(aes(group, mean_prop, fill = celltype)) +
  geom_col(position = "stack", width = 0.7) +
  scale_fill_manual(values = ct_pal) +
  labs(title = "Mean cell type proportions by group",
       x = NULL, y = "Mean proportion", fill = "Cell type") +
  theme_bw() +
  theme(axis.text.x  = element_text(angle = 45, hjust = 1),
        plot.title   = element_text(hjust = 0.5, face = "bold"),
        legend.text  = element_text(size = 7),
        legend.key.size = unit(0.4, "cm"))

ggsave(file.path(DIRS$s09, "stacked_bar_mean_proportion.pdf"),
       p_bar, width = 9, height = 6)
log_save(file.path(DIRS$s09, "stacked_bar_mean_proportion.pdf"))

# 9b — box plots: top 10 cell types by mean proportion, BPA vs Ctrl
top_ct <- prop_summary %>%
  group_by(celltype) %>%
  summarise(overall_mean = mean(mean_prop)) %>%
  slice_max(overall_mean, n = 10) %>%
  pull(celltype)

p_box <- props_long %>%
  filter(celltype %in% top_ct) %>%
  mutate(celltype = factor(celltype, levels = top_ct)) %>%
  ggplot(aes(treatment, proportion, fill = treatment)) +
  geom_boxplot(outlier.size = 0.8) +
  facet_grid(celltype ~ sex, scales = "free_y") +
  scale_fill_manual(values = c(Ctrl = "grey75",
                                BPA10ug = "#4DBBD5",
                                BPA10mg = "#E64B35")) +
  labs(title = "Top 10 cell types: BPA vs Ctrl",
       x = NULL, y = "Proportion", fill = "Treatment") +
  theme_bw(base_size = 9) +
  theme(axis.text.x  = element_text(angle = 45, hjust = 1),
        plot.title   = element_text(hjust = 0.5, face = "bold"),
        strip.text   = element_text(size = 7))

ggsave(file.path(DIRS$s09, "boxplot_top10_celltypes.pdf"),
       p_box, width = 8, height = 14)
log_save(file.path(DIRS$s09, "boxplot_top10_celltypes.pdf"))

# 9c — volcano: log2FC vs -log10(p_adj) for each comparison
p_volc <- stat_results %>%
  mutate(sig = p_adj < 0.05,
         label = ifelse(sig, celltype, NA)) %>%
  ggplot(aes(log2FC, -log10(p_adj), colour = sig, label = label)) +
  geom_point(size = 1.5) +
  geom_hline(yintercept = -log10(0.05), linetype = "dashed", colour = "grey50") +
  geom_vline(xintercept = 0, linetype = "dashed", colour = "grey50") +
  ggrepel::geom_text_repel(size = 2.5, max.overlaps = 15, na.rm = TRUE) +
  scale_colour_manual(values = c("FALSE" = "grey70", "TRUE" = "#E64B35")) +
  facet_grid(sex ~ comparison) +
  labs(title = "BPA effect on cell type proportions",
       x = "log2 FC (BPA / Ctrl)", y = "-log10(BH-adj p)",
       colour = "p_adj < 0.05") +
  theme_bw(base_size = 9) +
  theme(plot.title = element_text(hjust = 0.5, face = "bold"))

ggsave(file.path(DIRS$s09, "volcano_BPA_vs_Ctrl.pdf"),
       p_volc, width = 10, height = 7)
log_save(file.path(DIRS$s09, "volcano_BPA_vs_Ctrl.pdf"))

log_msg("=== All steps done ===")
cat(sprintf("\nAll results in: %s\n", RESULTS_DIR))