#!/usr/bin/env python
"""伪批量(pseudobulk)差异表达 vs 逐细胞检验 —— 用真实的多受试者公开数据当场对比。

数据:GEO GSE96583(Kang et al. 2018),8 位供者的 PBMC,每人各有对照与干扰素-β 刺激两份 ——
这是极少数「同一批人两种处理」的公开单细胞数据,天然适合配对设计。

要证明的事(我们在指南里反复强调、也是审稿人最常挑的一条):
把成千上万个细胞当独立样本做逐细胞检验,会把假阳性放大到荒唐的程度;
真正的样本量是**人数**,所以应该先按 供者×条件 把计数加总成伪批量,再用 DESeq2 这类方法检验。
本脚本两种做法都真跑一遍,把差距摆出来,并用已知的干扰素诱导基因(ISG)做正确性自检。
"""
import gzip, 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 scipy.io as sio
import scipy.sparse as sp

HERE = os.path.dirname(os.path.abspath(__file__))
D = os.path.join(HERE, "kang")
OUT = os.path.join(HERE, "out_pseudobulk")
os.makedirs(OUT, exist_ok=True)
M = {"dataset": "GEO GSE96583 (Kang et al. 2018) — 8 donors, control vs IFN-β"}

genes = [l.rstrip("\n").split("\t")[-1] for l in gzip.open(f"{D}/genes.tsv.gz", "rt")]
meta = pd.read_csv(f"{D}/meta.tsv.gz", sep="\t", header=0,
                   names=["tsne1", "tsne2", "ind", "stim", "cluster", "cell", "multiplets"])
meta.index.name = "barcode"

def load(tag, mtx, bcs):
    X = sio.mmread(f"{D}/{mtx}").tocsr()           # genes × cells
    bc = [l.strip() for l in gzip.open(f"{D}/{bcs}", "rt")]
    sub = meta[meta.stim == tag]
    idx = [i for i, b in enumerate(bc) if b in sub.index]
    X = X[:, idx].T.tocsr()                        # → cells × genes
    m = sub.loc[[bc[i] for i in idx]].copy()
    m["condition"] = tag
    return X, m

Xc, mc = load("ctrl", "ctrl.mtx.gz", "ctrl.barcodes.tsv.gz")
Xs, ms = load("stim", "stim.mtx.gz", "stim.barcodes.tsv.gz")
X = sp.vstack([Xc, Xs]).tocsr()
md = pd.concat([mc, ms])
M["cells_total"] = int(X.shape[0]); M["genes_total"] = len(genes)

keep = (md.multiplets == "singlet") & md.cell.notna() & (md.cell != "NA")
X, md = X[keep.values], md[keep]
M["cells_singlet"] = int(X.shape[0])
M["donors"] = sorted(md.ind.astype(str).unique().tolist())
M["cell_types"] = md.cell.value_counts().to_dict()

CT = "CD14+ Monocytes"                              # 干扰素反应最强的一群,信号明确
sel = (md.cell == CT).values
Xct, mdct = X[sel], md[sel]
M["focus_cell_type"] = CT
M["cells_in_focus"] = int(Xct.shape[0])
M["design"] = {"donors": len(M["donors"]), "conditions": 2,
               "true_n_per_group": len(M["donors"]), "cells_per_group": int(Xct.shape[0] / 2)}

# ---------- ① 正确做法:按 供者×条件 加总成伪批量,再上 DESeq2(配对设计) ----------
groups = mdct.groupby([mdct.ind.astype(str), mdct.condition]).indices
pb, rows = [], []
for (ind, cond), idx in groups.items():
    if len(idx) < 10:                               # 细胞太少的组合不进伪批量,并如实记录
        continue
    pb.append(np.asarray(Xct[idx].sum(axis=0)).ravel())
    rows.append({"donor": ind, "condition": cond, "cells": len(idx)})
pbdf = pd.DataFrame(np.vstack(pb), columns=genes,
                    index=[f"{r['donor']}_{r['condition']}" for r in rows])
pbdf = pbdf.loc[:, pbdf.sum(axis=0) >= 10]
pbdf = pbdf.loc[:, ~pbdf.columns.duplicated()]
clin = pd.DataFrame(rows, index=pbdf.index)
M["pseudobulk_samples"] = clin.to_dict("records")

from pydeseq2.dds import DeseqDataSet
from pydeseq2.ds import DeseqStats
dds = DeseqDataSet(counts=pbdf.astype(int), metadata=clin, design="~donor + condition", quiet=True)
dds.deseq2()
st = DeseqStats(dds, contrast=["condition", "stim", "ctrl"], quiet=True)
st.summary()
pbres = st.results_df.dropna(subset=["padj"])
pb_sig = pbres[(pbres.padj < 0.05) & (pbres.log2FoldChange.abs() > 1)]
M["pseudobulk_significant"] = int(len(pb_sig))
M["pseudobulk_tested"] = int(len(pbres))

# ---------- ② 常见错误做法:把每个细胞当独立样本做 Wilcoxon ----------
import scanpy as sc
import anndata as ad
A = ad.AnnData(X=Xct.astype(np.float32), obs=mdct.reset_index()[["condition", "ind"]].astype(str),
               var=pd.DataFrame(index=pd.Index(genes, name="gene").to_series().reset_index(drop=True)))
A.var_names = pd.Index(genes)
A.var_names_make_unique()
sc.pp.filter_genes(A, min_cells=10)
sc.pp.normalize_total(A, target_sum=1e4); sc.pp.log1p(A)
sc.tl.rank_genes_groups(A, "condition", groups=["stim"], reference="ctrl", method="wilcoxon")
cell = sc.get.rank_genes_groups_df(A, group="stim").dropna(subset=["pvals_adj"])
cell_sig = cell[(cell.pvals_adj < 0.05) & (cell.logfoldchanges.abs() > 1)]
M["percell_significant"] = int(len(cell_sig))
M["percell_tested"] = int(len(cell))
M["inflation_ratio"] = round(len(cell_sig) / max(len(pb_sig), 1), 1)

# ---------- ③ 正确性自检:已知的干扰素诱导基因两种方法都应排在最前 ----------
# 正确性自检:干扰素刺激后,这 10 个已知的 ISG 必须显著上调 —— 两种方法都查,逐个列出来
ISG = ["ISG15", "IFI6", "IFIT1", "IFIT3", "MX1", "OAS1", "IFI44L", "RSAD2", "CXCL10", "STAT1"]
cell_idx = cell.set_index("names")
isg_rows = []
for g in ISG:
    row = {"gene": g}
    if g in pbres.index:
        r = pbres.loc[g]
        row["pb_log2FC"] = round(float(r.log2FoldChange), 2)
        row["pb_padj"] = float(f"{r.padj:.2g}")
        row["pb_significant"] = bool(r.padj < 0.05 and abs(r.log2FoldChange) > 1)
    if g in cell_idx.index:
        c = cell_idx.loc[g]
        row["cell_log2FC"] = round(float(c.logfoldchanges), 2)
        row["cell_padj"] = float(f"{c.pvals_adj:.2g}")
    isg_rows.append(row)
M["isg_check"] = isg_rows
M["isg_significant_in_pseudobulk"] = sum(1 for r in isg_rows if r.get("pb_significant"))
M["top_pseudobulk"] = [{"gene": str(i), "log2FC": round(float(r.log2FoldChange), 2),
                        "padj": float(f"{r.padj:.3g}")}
                       for i, r in pbres.sort_values("padj").head(12).iterrows()]
overlap = set(pb_sig.index) & set(cell_sig.names)
M["overlap"] = int(len(overlap))
M["percell_only"] = int(len(set(cell_sig.names) - set(pb_sig.index)))
M["pseudobulk_only"] = int(len(set(pb_sig.index) - set(cell_sig.names)))
M["disagreement_pct"] = round(100 * (M["percell_only"] + M["pseudobulk_only"]) /
                              max(len(set(pb_sig.index) | set(cell_sig.names)), 1), 1)
pbres.to_csv(f"{OUT}/pseudobulk_results.csv")
cell.to_csv(f"{OUT}/percell_results.csv", index=False)

# ---------- 图 ----------
fig, ax = plt.subplots(figsize=(5.6, 4.4))
ax.bar(["Pseudobulk\n(n = %d donors)" % len(M["donors"]), "Per-cell test\n(n = %s cells)" % f"{M['cells_in_focus']:,}"],
       [M["pseudobulk_significant"], M["percell_significant"]], color=["#0d9488", "#dc2626"])
for i, v in enumerate([M["pseudobulk_significant"], M["percell_significant"]]):
    ax.text(i, v, f" {v:,}", ha="center", va="bottom", fontsize=11, fontweight="bold")
ax.set_ylabel("Genes called significant (padj<0.05, |log2FC|>1)")
ax.set_title("Same cells, same comparison, two methods", fontsize=11)
plt.tight_layout(); plt.savefig(f"{OUT}/method_comparison.png", dpi=140); plt.close()

pbres = pbres.assign(nl=-np.log10(pbres.padj.clip(lower=1e-300)))
sig = (pbres.padj < 0.05) & (pbres.log2FoldChange.abs() > 1)
plt.figure(figsize=(5.6, 4.6))
plt.scatter(pbres.log2FoldChange[~sig], pbres.nl[~sig], s=4, c="#b8c2d4", alpha=.5, edgecolors="none")
plt.scatter(pbres.log2FoldChange[sig], pbres.nl[sig], s=6, c="#0d9488", alpha=.8, edgecolors="none")
for g in [r["gene"] for r in M["isg_check"] if r.get("pb_significant")][:6]:
    if g in pbres.index:
        r = pbres.loc[g]
        plt.annotate(g, (r.log2FoldChange, r.nl), fontsize=8, color="#0f172a")
plt.axhline(-np.log10(0.05), c="#94a3b8", lw=.6, ls="--")
plt.xlabel("log2 fold change (IFN-β vs control)"); plt.ylabel("-log10 adjusted p")
plt.title(f"Pseudobulk DESeq2 — {CT}", fontsize=11)
plt.tight_layout(); plt.savefig(f"{OUT}/pseudobulk_volcano.png", dpi=140); 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_singlet", "focus_cell_type", "cells_in_focus", "design",
                                     "pseudobulk_significant", "percell_significant",
                                     "overlap", "percell_only", "pseudobulk_only", "disagreement_pct",
                                     "isg_significant_in_pseudobulk")}, indent=2))
for r in M["isg_check"]:
    print("  ", r)