# Cell Omics Guide / Shengxin Manmanxue: teaching adaptation, not a validated reproduction.
# Run one stage at a time from a new working directory. See the bilingual site guide.
# Prerequisites: Seurat 5, SeuratData + installed stxBrain, future, ggplot2,
# patchwork, sctransform, ape; harmony only for the multi stage.
# Download allen_cortex.rds using the official Seurat spatial vignette.
args <- commandArgs(trailingOnly = TRUE)
stage <- if (length(args)) args[[1]] else "help"
allowed <- c("audit", "domains", "cortex", "transfer", "multi-audit", "multi-domains", "multi-cortex", "multi")
if (!stage %in% allowed) stop(paste("Choose a stage:", paste(allowed, collapse = ", ")))
suppressPackageStartupMessages({ library(Seurat); library(ggplot2); library(patchwork) })
stopifnot(packageVersion("Seurat") >= "5.0.0")
future::plan(future::sequential)
set.seed(1224)
out <- "spatial-output"
dir.create(out, showWarnings = FALSE)
posterior <- stage %in% c("multi-audit", "multi-domains", "multi-cortex")
work <- if (posterior) file.path(out, "posterior-output") else out
dir.create(work, showWarnings = FALSE)
slice <- if (posterior) "posterior1" else "anterior1"
action <- sub("^multi-", "", stage)
p <- function(name) file.path(work, name)
put <- function(x, name) write.csv(x, p(name), row.names = FALSE)
save_plot <- function(x, name, width = 12, height = 7) ggsave(p(name), x, width = width, height = height, limitsize = FALSE)
load_object <- function(name) {
  stopifnot(file.exists(p(name)))
  readRDS(p(name))
}
review <- function(file, cells) {
  d <- read.csv(p(file), colClasses = "character", check.names = FALSE)
  stopifnot(all(c("barcode", "keep", "reason") %in% names(d)), !anyDuplicated(d$barcode), setequal(d$barcode, cells))
  d <- d[match(cells, d$barcode), , drop = FALSE]
  d$keep <- toupper(trimws(d$keep))
  if (anyNA(d) || any(!d$keep %in% c("TRUE", "FALSE")) || any(!nzchar(trimws(d$reason)))) {
    stop(paste("Review every keep (TRUE/FALSE) and reason in", p(file)))
  }
  d$keep <- d$keep == "TRUE"
  d
}
fit <- function(x) {
  if (ncol(x) < 60) stop("At least 60 reviewed spots are required by this 30-PC teaching configuration.")
  DefaultAssay(x) <- "Spatial"
  x <- SCTransform(x, assay = "Spatial", ncells = min(3000, ncol(x)), verbose = FALSE)
  if (length(VariableFeatures(x)) <= 30) stop("Too few variable features for 30 PCs.")
  RunPCA(x, assay = "SCT", npcs = 30, verbose = FALSE)
}
map_labels <- function(query, reference) {
  DefaultAssay(query) <- "SCT"
  DefaultAssay(reference) <- "SCT"
  shared <- intersect(rownames(reference), rownames(query))
  if (length(shared) < 200) stop("Fewer than 200 shared genes: audit species and feature naming.")
  anchors <- FindTransferAnchors(reference = reference, query = query,
    normalization.method = "SCT", reference.reduction = "pca", dims = 1:30)
  scores <- TransferData(anchorset = anchors, refdata = reference$subclass,
    prediction.assay = TRUE, weight.reduction = query[["pca"]], dims = 1:30)
  query[["predictions"]] <- scores
  m <- GetAssayData(query, assay = "predictions", layer = "data")
  m <- m[setdiff(rownames(m), "max"), colnames(query), drop = FALSE]
  stopifnot(nrow(m) > 1, all(is.finite(m)))
  # Small label-by-spot matrix only; never densify the gene-by-cell counts.
  small <- as.matrix(m)
  best <- max.col(t(small), ties.method = "first")
  query$prediction_max <- apply(small, 2, max)
  query$prediction_label <- ifelse(query$prediction_max < 0.5, "uncertain", rownames(small)[best])
  list(object = query, shared = shared, scores = data.frame(barcode = colnames(query),
    label = query$prediction_label, max_score = query$prediction_max, t(small), check.names = FALSE))
}

if (action == "audit") {
  brain <- SeuratData::LoadData("stxBrain", type = slice)
  stopifnot(length(Images(brain)) == 1, !anyDuplicated(colnames(brain)))
  coords <- GetTissueCoordinates(brain)
  stopifnot(all(colnames(brain) %in% rownames(coords)))
  saveRDS(brain, p("brain-raw.rds"))
  capture.output(list(slice = slice, dimensions = dim(brain), images = Images(brain), coordinates = head(coords)), file = p("input-audit.txt"))
  save_plot(VlnPlot(brain, features = c("nCount_Spatial", "nFeature_Spatial"), pt.size = 0) /
    SpatialFeaturePlot(brain, features = "nCount_Spatial"), "tissue-QC.pdf")
  if (!file.exists(p("qc-review.csv"))) put(data.frame(barcode = colnames(brain),
    nCount = brain$nCount_Spatial, nFeature = brain$nFeature_Spatial, keep = "", reason = ""), "qc-review.csv")
}
if (action == "domains") {
  brain <- load_object("brain-raw.rds")
  decisions <- review("qc-review.csv", colnames(brain))
  brain <- subset(brain, cells = decisions$barcode[decisions$keep])
  brain <- fit(brain)
  brain <- FindNeighbors(brain, reduction = "pca", dims = 1:30)
  brain <- FindClusters(brain, resolution = 0.5, random.seed = 1224)
  brain <- RunUMAP(brain, dims = 1:30, seed.use = 1224)
  saveRDS(brain, p("brain.rds"))
  put(decisions, "qc-decisions.csv")
  markers <- intersect(c("Hpca", "Ttr", "Slc17a7", "Cux2", "Rorb", "Tle4", "Gad1", "Plp1"), rownames(brain))
  stopifnot(length(markers) > 0)
  save_plot((DimPlot(brain, label = TRUE) | SpatialDimPlot(brain, label = TRUE)) /
    SpatialFeaturePlot(brain, features = markers, ncol = 4), "domains.pdf", height = 12)
  if (!file.exists(p("cortex-review.csv"))) put(data.frame(barcode = colnames(brain),
    cluster = as.character(Idents(brain)), keep = "", reason = ""), "cortex-review.csv")
}
if (action == "cortex") {
  brain <- load_object("brain.rds")
  decisions <- review("cortex-review.csv", colnames(brain))
  cortex <- subset(brain, cells = decisions$barcode[decisions$keep])
  cortex <- fit(cortex)
  features <- head(VariableFeatures(cortex), 200)
  cortex <- FindSpatiallyVariableFeatures(cortex, assay = "SCT", features = features,
    image = Images(cortex)[[1]], selection.method = "moransi")
  put(data.frame(feature = rownames(cortex[["SCT"]][[]]), cortex[["SCT"]][[]], check.names = FALSE), "moran-features.csv")
  writeLines(features, p("moran-tested.txt"))
  saveRDS(cortex, p("cortex.rds"))
  put(decisions, "cortex-decisions.csv")
  save_plot(SpatialDimPlot(cortex, crop = FALSE), "cortex.pdf")
}
if (action == "transfer") {
  cortex <- load_object("cortex.rds")
  stopifnot(file.exists("allen_cortex.rds"))
  reference <- readRDS("allen_cortex.rds")
  stopifnot(inherits(reference, "Seurat"), "subclass" %in% colnames(reference[[]]),
    !anyNA(reference$subclass), all(nzchar(as.character(reference$subclass))), "RNA" %in% Assays(reference))
  reference <- SCTransform(reference, assay = "RNA", ncells = min(3000, ncol(reference)), verbose = FALSE)
  reference <- RunPCA(reference, npcs = 30, verbose = FALSE)
  saveRDS(reference, p("reference-sct.rds"))
  result <- map_labels(cortex, reference)
  saveRDS(result$object, p("cortex-predictions.rds"))
  put(result$scores, "prediction-scores.csv")
  writeLines(result$shared, p("shared-genes.txt"))
  save_plot(SpatialDimPlot(result$object, group.by = "prediction_label") |
    SpatialFeaturePlot(result$object, features = "prediction_max"), "predictions.pdf")
}
if (stage == "multi") {
  stopifnot(requireNamespace("harmony", quietly = TRUE))
  a <- load_object("cortex-predictions.rds")
  b <- readRDS(file.path(out, "posterior-output", "cortex.rds"))
  reference <- load_object("reference-sct.rds")
  result <- map_labels(b, reference)
  b <- result$object
  saveRDS(b, p("posterior-predictions.rds"))
  put(result$scores, "posterior-prediction-scores.csv")
  save_plot(SpatialDimPlot(b, group.by = "prediction_label") |
    SpatialFeaturePlot(b, features = "prediction_max"), "posterior-predictions.pdf")
  a$sample_id <- "anterior1"
  b$sample_id <- "posterior1"
  DefaultAssay(a) <- DefaultAssay(b) <- "SCT"
  features <- intersect(VariableFeatures(a), VariableFeatures(b))
  if (length(features) < 200) stop("Insufficient shared variable features; inspect feature selection.")
  st <- merge(a, y = b, add.cell.ids = c("anterior1", "posterior1"), merge.data = TRUE)
  # Recreate shared residual scale data from per-slice SCT models, without dense counts.
  parts <- PrepSCTIntegration(list(a, b), anchor.features = features, verbose = FALSE)
  scales <- lapply(parts, function(x) GetAssayData(x, assay = "SCT", layer = "scale.data")[features, , drop = FALSE])
  colnames(scales[[1]]) <- paste0("anterior1_", colnames(scales[[1]]))
  colnames(scales[[2]]) <- paste0("posterior1_", colnames(scales[[2]]))
  scale <- do.call(cbind, scales)
  st <- SetAssayData(st, assay = "SCT", layer = "scale.data", new.data = scale[, colnames(st), drop = FALSE])
  st <- RunPCA(st, assay = "SCT", features = features, npcs = 30, verbose = FALSE)
  st <- RunUMAP(st, reduction = "pca", dims = 1:30, reduction.name = "umap.pca", reduction.key = "uPCA_", seed.use = 1224)
  st <- harmony::RunHarmony(st, group.by.vars = "sample_id", reduction.use = "pca", dims.use = 1:30)
  st <- RunUMAP(st, reduction = "harmony", dims = 1:30, reduction.name = "umap.harmony", reduction.key = "uHarmony_", seed.use = 1224)
  saveRDS(st, p("multislice.rds"))
  save_plot((DimPlot(st, reduction = "umap.pca", group.by = "sample_id") |
    DimPlot(st, reduction = "umap.harmony", group.by = "sample_id")) /
    SpatialDimPlot(st, group.by = "prediction_label", ncol = 2), "integration-comparison.pdf", height = 12)
}
capture.output(sessionInfo(), file = p(paste0("session-", stage, ".txt")))
message("Stage completed. Review artifacts in ", work, " before marking the stage verified.")