#!/usr/bin/env python
"""影像式空间转录组(Xenium)样例。公开数据:10x Xenium 人乳腺癌 FFPE 切片 Rep1。

Visium 那套在这里全部不适用,而且这不是「差不多就行」的差别:

  · Visium 一个 spot 覆盖多个细胞,Xenium 是**单细胞分辨率** —— 但代价是只测面板里的基因;
  · 因此**面板就是天花板**:313 个基因之外的东西不是「没检测到」,是根本没被问过。
    拿它去做「发现新细胞类型」是方法层面的错配,不是分析水平问题;
  · 每个细胞的转录本数量比 scRNA 低一到两个数量级(这里中位 164 个),
    照搬 scRNA 的「每细胞 ≥500 基因」阈值会把整张切片砍光;
  · 多出来一个 scRNA 完全没有的失败模式:**细胞分割**。转录本被算到哪个细胞头上,
    是算法画出来的边界决定的;边界画错,表达谱就是两个细胞的混合。

这个脚本只回答三个能被验证的问题:
  ① 背景有多高 —— 阴性对照探针/码字占多少信号(这是整张片子的信噪比下限,和 CITE-seq 的同型对照同理);
  ② 分割质量如何 —— 用细胞核占细胞面积的比例、以及「只有核没有胞质」的细胞比例来量;
  ③ 面板够不够 —— 已知乳腺谱系标志基因在不在面板里、能不能把群分开。
"""
import json, os, warnings
warnings.filterwarnings("ignore")
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import scanpy as sc

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(HERE, "out_xenium")
os.makedirs(OUT, exist_ok=True)
sc.settings.figdir = OUT
sc.settings.set_figure_params(dpi=140, frameon=False, figsize=(5.0, 4.2), facecolor="white")
M = {"dataset": "10x Genomics Xenium human breast cancer FFPE, replicate 1 (public)"}

adata = sc.read_10x_mtx(os.path.join(HERE, "data", "xenium_cfm", "cell_feature_matrix"),
                        var_names="gene_symbols", make_unique=True)
feat = pd.read_csv(os.path.join(HERE, "data", "xenium_cfm", "cell_feature_matrix", "features.tsv.gz"),
                   sep="\t", header=None, names=["id", "name", "type"])
M["cells_raw"] = int(adata.n_obs)
M["panel_genes"] = int((feat["type"] == "Gene Expression").sum())
M["control_features"] = {t: int(n) for t, n in feat["type"].value_counts().items() if t != "Gene Expression"}

cells = pd.read_csv(os.path.join(HERE, "data", "xenium_cells.csv.gz"))
# ⚠️ 矩阵里的 barcode 是**字符串** '1','2','3',而 cells.csv 里的 cell_id 是**整数** 1,2,3。
# 直接按原类型对齐,交集为空 —— 而 pandas 的 reindex **不会报错**,它安静地给你一整列 NaN。
# 第一版就是这样:合并「成功」,12 分钟后才在画直方图时炸出来。要是那一步宽容一点,
# 这些 NaN 会一路进到 metrics.json 和对外页面上。所以:两边都转成字符串,并且**断言真的对上了**。
cells["cell_id"] = cells["cell_id"].astype(str)
adata.obs_names = adata.obs_names.astype(str)
overlap = len(set(adata.obs_names) & set(cells["cell_id"]))
if overlap < 0.99 * adata.n_obs:
    raise SystemExit(f"细胞元数据对不上:矩阵 {adata.n_obs} 个细胞,只匹配上 {overlap} 个 —— 先查 ID 格式,不要继续算")
M["metadata_join_matched"] = overlap
cells = cells.set_index("cell_id").reindex(adata.obs_names)
for c in ("transcript_counts", "control_probe_counts", "control_codeword_counts",
          "total_counts", "cell_area", "nucleus_area", "x_centroid", "y_centroid"):
    v = cells[c].values.astype(float)
    if not np.isfinite(v).all():
        raise SystemExit(f"列 {c} 里有 NaN —— 合并出了问题,停下来查,别拿 NaN 往下算")
    adata.obs[c] = v

# ---------- ① 背景:阴性对照占多少 ----------
ctrl = adata.obs["control_probe_counts"].values + adata.obs["control_codeword_counts"].values
tot = np.maximum(adata.obs["total_counts"].values, 1)
M["median_transcripts_per_cell"] = float(np.median(adata.obs["transcript_counts"]))
M["control_share_pct"] = round(100 * float(ctrl.sum() / tot.sum()), 3)
M["cells_with_any_control"] = int((ctrl > 0).sum())
M["cells_with_any_control_pct"] = round(100 * float((ctrl > 0).mean()), 1)
# 面板里真基因的平均每细胞计数 vs 单个阴性对照探针的平均计数 —— 信噪比的直观说法
n_ctrl_feat = sum(M["control_features"].values())
M["signal_to_background_ratio"] = round(
    float((adata.obs["transcript_counts"].mean() / M["panel_genes"]) /
          max(ctrl.mean() / max(n_ctrl_feat, 1), 1e-9)), 1)

# ---------- ② 分割质量 ----------
area = adata.obs["cell_area"].values
nuc = adata.obs["nucleus_area"].values
ratio = nuc / np.maximum(area, 1e-9)
adata.obs["nucleus_fraction"] = ratio
M["median_cell_area_um2"] = round(float(np.median(area)), 1)
M["median_nucleus_fraction"] = round(float(np.median(ratio)), 3)
# 「核几乎等于整个细胞」= 胞质没被分进来,这类细胞的表达谱偏向核内转录本
M["nucleus_only_cells_pct"] = round(100 * float((ratio > 0.9).mean()), 1)
# 面积异常大的,常见于两个细胞被合成一个
big = np.percentile(area, 99.5)
M["oversized_cells_pct"] = round(100 * float((area > big).mean()), 2)
M["empty_cells"] = int((adata.obs["transcript_counts"].values == 0).sum())

fig, axes = plt.subplots(1, 3, figsize=(12.6, 3.5))
axes[0].hist(adata.obs["transcript_counts"], bins=80, color="#0a6ed1", range=(0, 800))
axes[0].set_title(f"Transcripts per cell (median {M['median_transcripts_per_cell']:.0f})", fontsize=10)
axes[1].hist(ratio, bins=80, color="#0d9488")
axes[1].axvline(0.9, color="#dc2626", ls="--", lw=1.2)
axes[1].set_title("Nucleus area / cell area", fontsize=10)
axes[2].hist(np.log10(np.maximum(area, 1)), bins=80, color="#7c3aed")
axes[2].set_title("log10 cell area (µm²)", fontsize=10)
for a in axes:
    a.set_yticks([])
plt.tight_layout(); plt.savefig(f"{OUT}/xenium_qc.png", dpi=140, bbox_inches="tight"); plt.close()

# ---------- 过滤 + 聚类(阈值必须按这个平台定,不能照搬 scRNA)----------
keep = (adata.obs["transcript_counts"].values >= 10) & (area > 0)
M["min_transcripts_threshold"] = 10
M["cells_after_qc"] = int(keep.sum())
M["removed_pct"] = round(100 * (1 - keep.sum() / adata.n_obs), 1)
ad = adata[keep].copy()
ad.layers["counts"] = ad.X.copy()
sc.pp.normalize_total(ad)
sc.pp.log1p(ad)
sc.pp.pca(ad, n_comps=30)
sc.pp.neighbors(ad, n_neighbors=15)
sc.tl.leiden(ad, resolution=0.6, key_added="cluster", flavor="igraph", n_iterations=2, directed=False)
M["clusters"] = int(ad.obs.cluster.nunique())

# ---------- ③ 面板够不够:已知乳腺谱系标志在不在、分不分得开 ----------
WANTED = {
    "Epithelial / luminal": ["EPCAM", "KRT8", "KRT18", "ESR1"],
    "Basal / myoepithelial": ["KRT14", "KRT5", "ACTA2", "TP63"],
    "T cells": ["CD3D", "CD3E", "CD8A", "IL7R"],
    "B / plasma": ["MS4A1", "CD79A", "JCHAIN"],
    "Myeloid": ["ITGAX", "CD68", "LYZ", "CD14"],
    "Endothelial": ["PECAM1", "VWF", "CDH5"],
    "Fibroblast": ["PDGFRA", "PDGFRB", "LUM", "DCN"],
}
panel = set(ad.var_names)
cover = {}
for k, genes in WANTED.items():
    present = [g for g in genes if g in panel]
    cover[k] = {"asked": len(genes), "in_panel": len(present), "present": present}
M["marker_coverage"] = cover
M["lineages_fully_missing"] = [k for k, v in cover.items() if v["in_panel"] == 0]
M["lineages_partially_covered"] = [k for k, v in cover.items() if 0 < v["in_panel"] < v["asked"]]

# 每个群里表达最高的面板基因 —— 用来看群是不是真的对应可解释的谱系
sc.tl.rank_genes_groups(ad, "cluster", method="wilcoxon", n_genes=5)
top = {c: [str(g) for g in ad.uns["rank_genes_groups"]["names"][c]] for c in ad.obs.cluster.cat.categories}
M["top_markers_per_cluster"] = top
found_lineages = 0
for k, v in cover.items():
    if v["in_panel"] and any(any(g in row for g in v["present"]) for row in top.values()):
        found_lineages += 1
M["lineages_recovered_in_clusters"] = found_lineages
M["lineages_checked"] = len(WANTED)
M["sanity_check_passed"] = bool(found_lineages >= len(WANTED) - 2 and M["control_share_pct"] < 1)

sc.pl.embedding(ad, basis="spatial", color="cluster", save=None, show=False) if "spatial" in ad.obsm else None
fig, ax = plt.subplots(figsize=(6.0, 6.0))
cl = ad.obs["cluster"].astype(str).values
codes, uniq = pd.factorize(cl)
sc_ = ax.scatter(ad.obs["x_centroid"], -ad.obs["y_centroid"], c=codes, s=0.35, cmap="tab20", linewidths=0)
ax.set_aspect("equal"); ax.set_axis_off()
ax.set_title(f"{M['clusters']} clusters in tissue coordinates", fontsize=11)
plt.tight_layout(); plt.savefig(f"{OUT}/xenium_map.png", dpi=150, bbox_inches="tight"); plt.close()

with open(f"{OUT}/metrics.json", "w") as f:
    json.dump(M, f, indent=2, ensure_ascii=False)
print(json.dumps({k: M[k] for k in (
    "cells_raw", "panel_genes", "control_features", "median_transcripts_per_cell",
    "control_share_pct", "cells_with_any_control_pct", "signal_to_background_ratio",
    "median_cell_area_um2", "median_nucleus_fraction", "nucleus_only_cells_pct",
    "oversized_cells_pct", "cells_after_qc", "removed_pct", "clusters",
    "lineages_fully_missing", "lineages_recovered_in_clusters", "sanity_check_passed")},
    indent=2, ensure_ascii=False))