#!/usr/bin/env python
"""Multiome(同一个细胞里同时测 RNA 和 ATAC)样例。公开数据:10x PBMC granulocyte-sorted 3k。

多组学最常被拿来讲故事的一句话是「染色质开放 → 基因表达」。但这份数据的价值恰恰在于:
RNA 和 ATAC 是**在同一个细胞里**测的,所以「两者是否一致」第一次成了可以直接量的东西,
而不是靠两批实验对着猜。这个脚本只回答两个能被验证的问题:

  ① 同一批细胞,分别只用 RNA 和只用 ATAC 聚类,结果一致吗?
     —— 如果不一致,那么「这个细胞是什么类型」就取决于你用了哪一半数据,
        任何联合分析都必须说清是谁在主导。用 ARI 量化。
  ② 基因附近的染色质开放程度,能不能预测这个基因的表达量?
     —— 用文件自带的基因组坐标算「基因活性」(基因体 ±2kb 内的峰),再和 RNA 比。
        预期是正相关但远不完美:**开放是必要不充分条件**,这一条经常被写反。

自检:CD3D/MS4A1/LYZ 这类谱系标志基因,应当在对应细胞群里同时出现「表达高」和「附近开放」;
若连这个都对不上,说明坐标或配对错了,后面的结论一律不可信。
"""
import json, os, re, 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
import scipy.sparse as sp
from scipy.sparse.linalg import svds
from sklearn.metrics import adjusted_rand_score, normalized_mutual_info_score

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(HERE, "out_multiome")
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 PBMC granulocyte-sorted 3k Multiome (public)"}

full = sc.read_10x_h5(os.path.join(HERE, "data", "multiome_3k.h5"), gex_only=False)
full.var_names_make_unique()
is_gene = (full.var["feature_types"] == "Gene Expression").values
is_peak = (full.var["feature_types"] == "Peaks").values
M["cells_raw"] = int(full.n_obs)
M["genes"], M["peaks"] = int(is_gene.sum()), int(is_peak.sum())

# ---------- 质控:两个模态各有各的指标,必须分别看 ----------
per = pd.read_csv(os.path.join(HERE, "data", "multiome_3k_metrics.csv"), index_col=0)
per = per.loc[[b for b in full.obs_names if b in per.index]]
full = full[per.index].copy()
# ⚠️ 这份**官方公开数据**自带的 per_barcode_metrics.csv 里,gex_umis_count / gex_genes_count
# 整列全是 0(66 万行没有一个非零),而 gex_raw_reads 是有值的 —— 列在、名字对、类型对,就是没填。
# 任何按 `gex_umis_count >= 1000` 过滤的流程都会**静默丢掉 100% 的细胞**。
# 这里的做法是:凡是要拿来做过滤判据的列,先检查它是不是整列为零;是就拒绝使用,改从矩阵自己算。
def usable(col):
    if col not in per.columns:
        return None
    v = per[col].values.astype(float)
    return None if not np.isfinite(v).any() or (v == 0).all() else v

M["unusable_metric_columns"] = [c for c in ("gex_umis_count", "gex_genes_count", "atac_fragments")
                                if c in per.columns and usable(c) is None]
_gu, _gg = usable("gex_umis_count"), usable("gex_genes_count")
_gene_X = full[:, is_gene].X
gex_umi = _gu if _gu is not None else np.asarray(_gene_X.sum(1)).ravel().astype(float)
gex_genes = _gg if _gg is not None else np.asarray((_gene_X > 0).sum(1)).ravel().astype(float)
M["gex_metrics_recomputed_from_matrix"] = bool(_gu is None or _gg is None)
atac_frag = usable("atac_fragments")
if atac_frag is None:
    atac_frag = per["atac_raw_reads"].values.astype(float)
M["median_gex_umi"] = float(np.median(gex_umi))
M["median_gex_genes"] = float(np.median(gex_genes))
M["median_atac_fragments"] = float(np.median(atac_frag))

rna = full[:, is_gene].copy()
atac = full[:, is_peak].copy()
rna.var_names_make_unique()
sc.pp.calculate_qc_metrics(rna, qc_vars=[], inplace=True, percent_top=None, log1p=False)
mt = rna.var_names.str.startswith("MT-")
rna.obs["pct_mt"] = np.asarray(rna[:, mt].X.sum(1)).ravel() / np.maximum(np.asarray(rna.X.sum(1)).ravel(), 1) * 100

# 两边都要合格才留 —— multiome 的一个真实风险是「一半好一半坏」的细胞
keep = (gex_genes >= 500) & (gex_umi >= 1000) & (atac_frag >= 1000) & (rna.obs["pct_mt"].values < 20)
M["cells_after_qc"] = int(keep.sum())
M["removed_pct"] = round(100 * (1 - keep.sum() / M["cells_raw"]), 1)
M["removed_for_rna_only"] = int((~((gex_genes >= 500) & (gex_umi >= 1000)) & (atac_frag >= 1000)).sum())
M["removed_for_atac_only"] = int((((gex_genes >= 500) & (gex_umi >= 1000)) & ~(atac_frag >= 1000)).sum())
rna, atac = rna[keep].copy(), atac[keep].copy()

# ---------- ① 只用 RNA 聚类 ----------
rna.layers["counts"] = rna.X.copy()
sc.pp.normalize_total(rna, target_sum=1e4)
sc.pp.log1p(rna)
sc.pp.highly_variable_genes(rna, n_top_genes=2000)
rna_hvg = rna[:, rna.var.highly_variable].copy()
sc.pp.scale(rna_hvg, max_value=10)
sc.tl.pca(rna_hvg, n_comps=30)
sc.pp.neighbors(rna_hvg, n_neighbors=15)
sc.tl.umap(rna_hvg)
sc.tl.leiden(rna_hvg, resolution=0.6, key_added="rna_clusters",
             flavor="igraph", n_iterations=2, directed=False)
M["rna_clusters"] = int(rna_hvg.obs.rna_clusters.nunique())

# ---------- ② 只用 ATAC 聚类(TF-IDF + LSI,丢掉深度成分)----------
X = (atac.X > 0).astype(np.float32)
tf = X.multiply(sp.csr_matrix(1.0 / np.maximum(np.asarray(X.sum(1)), 1)))
idf = X.shape[0] / np.maximum(np.asarray(X.sum(0)).ravel(), 1)
tfidf = tf.multiply(sp.csr_matrix(idf))
tfidf = (tfidf * 1e4).tocsr()
tfidf.data = np.log1p(tfidf.data)
U, S, _ = svds(tfidf.astype(np.float64), k=30)
order = np.argsort(S)[::-1]
lsi = U[:, order] * S[order]
depth = np.log10(np.maximum(atac_frag[keep], 1))
M["lsi1_depth_correlation"] = round(abs(float(np.corrcoef(lsi[:, 0], depth)[0, 1])), 3)
atac.obsm["X_lsi"] = lsi[:, 1:]
sc.pp.neighbors(atac, use_rep="X_lsi", n_neighbors=15)
sc.tl.leiden(atac, resolution=0.6, key_added="atac_clusters",
             flavor="igraph", n_iterations=2, directed=False)
M["atac_clusters"] = int(atac.obs.atac_clusters.nunique())

# ---------- ③ 两个模态给出的分群一致吗 ----------
r_lab = rna_hvg.obs["rna_clusters"].astype(str).values
a_lab = atac.obs["atac_clusters"].astype(str).values
M["ARI_rna_vs_atac"] = round(float(adjusted_rand_score(r_lab, a_lab)), 3)
M["NMI_rna_vs_atac"] = round(float(normalized_mutual_info_score(r_lab, a_lab)), 3)
ct = pd.crosstab(pd.Series(r_lab, name="RNA"), pd.Series(a_lab, name="ATAC"))
# 每个 RNA 群里,最大的那个 ATAC 群占多少 —— 「一个 RNA 群是否被 ATAC 拆开」
purity = (ct.max(axis=1) / ct.sum(axis=1))
M["rna_cluster_purity_in_atac"] = {k: round(float(v), 3) for k, v in purity.items()}
M["median_purity"] = round(float(purity.median()), 3)
M["worst_purity_cluster"] = str(purity.idxmin())
M["worst_purity"] = round(float(purity.min()), 3)

fig, ax = plt.subplots(figsize=(6.2, 5.0))
im = ax.imshow(ct.values / ct.values.sum(axis=1, keepdims=True), cmap="Blues", aspect="auto")
ax.set_xlabel("ATAC-only clusters"); ax.set_ylabel("RNA-only clusters")
ax.set_xticks(range(ct.shape[1])); ax.set_xticklabels(ct.columns, fontsize=7)
ax.set_yticks(range(ct.shape[0])); ax.set_yticklabels(ct.index, fontsize=7)
ax.set_title(f"Same cells, clustered twice — ARI {M['ARI_rna_vs_atac']}", fontsize=11)
plt.colorbar(im, ax=ax, shrink=0.8, label="share of RNA cluster")
plt.tight_layout(); plt.savefig(f"{OUT}/modality_agreement.png", dpi=140, bbox_inches="tight"); plt.close()

rna_hvg.obs["atac_clusters"] = pd.Categorical(a_lab)
sc.pl.umap(rna_hvg, color=["rna_clusters", "atac_clusters"], ncols=2,
           title=["Clusters from RNA", "ATAC clusters, on the RNA map"],
           save="_multiome_clusters.png", show=False)

# ---------- ④ 基因附近开放 → 表达?用文件自带坐标算基因活性 ----------
def parse(iv):
    m = re.match(r"^(chr[\w]+):(\d+)-(\d+)$", str(iv))
    return (m.group(1), int(m.group(2)), int(m.group(3))) if m else (None, 0, 0)

pk = pd.DataFrame([parse(v) for v in atac.var["interval"]], columns=["chr", "start", "end"])
gn = pd.DataFrame([parse(v) for v in rna.var["interval"]], columns=["chr", "start", "end"])
gn["gene"] = rna.var_names
FLANK = 2000
Xp = (atac.X > 0).astype(np.float32).tocsc()
# ⚠️ 「基因活性」= 基因附近开放峰的个数,这个数**天然随细胞测序深度走**:
# 测得深的细胞到处都开,浅的细胞到处都不开。不做深度归一化就直接和表达比,
# 比的其实是深度而不是调控。第一版就是这么写的,自检当场报错:
# 十个谱系标志基因里只有 4 个的「表达最高群」和「开放最高群」对得上 —— 这不是生物学,是 bug。
_cell_open = np.asarray(Xp.sum(axis=1)).ravel()
_norm = np.median(_cell_open) / np.maximum(_cell_open, 1)
act, npeak = {}, {}
for chrom, gsub in gn.dropna(subset=["chr"]).groupby("chr"):
    psub = pk[pk["chr"] == chrom]
    if psub.empty:
        continue
    starts = psub["start"].values; ends = psub["end"].values; idx = psub.index.values
    o = np.argsort(starts); starts, ends, idx = starts[o], ends[o], idx[o]
    for _, g in gsub.iterrows():
        lo, hi = g["start"] - FLANK, g["end"] + FLANK
        j0 = np.searchsorted(starts, lo - 5000, side="left")
        j1 = np.searchsorted(starts, hi, side="right")
        if j1 <= j0:
            continue
        sel = idx[j0:j1][ends[j0:j1] >= lo]
        if len(sel):
            act[g["gene"]] = np.asarray(Xp[:, sel].sum(axis=1)).ravel() * _norm
            npeak[g["gene"]] = len(sel)
M["genes_with_nearby_peaks"] = len(act)

MARKERS = ["CD3D", "CD3E", "IL7R", "MS4A1", "CD79A", "LYZ", "CD14", "NKG7", "GNLY", "FCGR3A", "PPBP"]
have = [g for g in MARKERS if g in act and g in rna.var_names]
rna_log = rna.copy()
npeaks = {g: int(n) for g, n in npeak.items()}
from scipy.stats import spearmanr

# 自检的正确形式:不是「表达最高的那个群必须等于开放最高的那个群」。
# 第一版就是那么写的,十个标志基因只对上 3–4 个,看起来像大失败;
# 逐个查下去发现 CD3D 的表达高在 4/5/6/7 群、开放也高在 4/5/6/7 群 —— **全是 T 细胞群**,
# 只是群内排序差了一点点(0.281 vs 0.297)。生物学是对的,是判据太脆。
# 改成跨细胞群的秩相关:方向对不对,而不是最大值落在哪一格。
rows = []
for g in have:
    e = np.asarray(rna_log[:, g].X.todense()).ravel()
    df = pd.DataFrame({"expr": e, "act": act[g], "cl": r_lab}).groupby("cl").mean()
    rho = float(spearmanr(df["expr"], df["act"]).statistic)
    rows.append({"gene": g, "nearby_peaks": npeaks.get(g, 0),
                 "spearman_across_clusters": round(rho, 3) if rho == rho else None,
                 "agrees": bool(rho == rho and rho >= 0.5)})
M["marker_agreement"] = rows
M["markers_checked"] = len(rows)
M["markers_agreeing"] = int(sum(x["agrees"] for x in rows))

# 全局:表达 vs 附近开放的相关,**按附近峰数分层**。
# 这是这份样品真正值钱的结论 —— 「基因活性」被广泛用来给 ATAC 标细胞类型,
# 但它可靠不可靠,取决于这个基因附近到底有几个峰,而这一点几乎从不被报告。
buckets = {"1 peak": (1, 1), "2 peaks": (2, 2), "3-5 peaks": (3, 5), "6+ peaks": (6, 10**6)}
strat = {}
allr = []
for g in list(act.keys()):
    if g not in set(rna.var_names):
        continue
    e = np.asarray(rna_log[:, g].X.todense()).ravel()
    df = pd.DataFrame({"e": e, "a": act[g], "cl": r_lab}).groupby("cl").mean()
    if df["e"].std() == 0 or df["a"].std() == 0:
        continue
    r = float(np.corrcoef(df["e"], df["a"])[0, 1])
    if r != r:
        continue
    allr.append((npeaks.get(g, 0), r))
allr_arr = np.array(allr)
M["genes_correlated"] = int(len(allr_arr))
M["median_expr_vs_accessibility_r"] = round(float(np.median(allr_arr[:, 1])), 3)
for name, (lo, hi) in buckets.items():
    sel = allr_arr[(allr_arr[:, 0] >= lo) & (allr_arr[:, 0] <= hi)]
    if len(sel):
        strat[name] = {"genes": int(len(sel)), "median_r": round(float(np.median(sel[:, 1])), 3),
                       "pct_negative": round(100 * float((sel[:, 1] < 0).mean()), 1)}
M["by_peak_count"] = strat
M["pct_genes_negative_r"] = round(100 * float((allr_arr[:, 1] < 0).mean()), 1)
M["median_nearby_peaks"] = int(np.median(allr_arr[:, 0]))
M["sanity_check_passed"] = bool(M["markers_agreeing"] >= len(rows) - 1)

fig, ax = plt.subplots(figsize=(6.4, 3.6))
names = list(strat.keys())
ax.bar(names, [strat[n]["median_r"] for n in names], color="#0d9488")
for i, n in enumerate(names):
    ax.text(i, strat[n]["median_r"] + 0.01, f"n={strat[n]['genes']}", ha="center", fontsize=8)
ax.axhline(0, color="#111", lw=0.8)
ax.set_ylabel("median r (expression vs accessibility)")
ax.set_title("Gene activity is only as good as the peaks near the gene", fontsize=11)
plt.tight_layout(); plt.savefig(f"{OUT}/expr_vs_accessibility.png", dpi=140, 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", "cells_after_qc", "removed_for_rna_only", "removed_for_atac_only",
    "rna_clusters", "atac_clusters", "ARI_rna_vs_atac", "NMI_rna_vs_atac",
    "median_purity", "worst_purity_cluster", "worst_purity",
    "median_expr_vs_accessibility_r", "median_nearby_peaks", "pct_genes_negative_r",
    "by_peak_count", "markers_checked", "markers_agreeing", "sanity_check_passed")}, indent=2, ensure_ascii=False))