#!/usr/bin/env python
"""细胞组成(差异丰度)分析 —— 正确做法 vs 常见错误做法。

数据:GEO GSE96583(8 位供者,每人各有对照与 IFN-β 刺激),已在本地。

要证明的事:比较两组之间「某类细胞变多了没有」时,
① 把所有细胞倒在一起做卡方/Fisher —— 细胞数以万计,几乎任何微小差异都会显著,等于凭空造出组成变化;
② 正确做法是每位供者算一次比例,再在**供者层面**做配对检验(n=8,不是 n=24000)。
这份样品把两种做法并排跑出来,让人看见差多少。
"""
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
from scipy import stats

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

meta = pd.read_csv(f"{D}/meta.tsv.gz", sep="\t", header=0,
                   names=["tsne1", "tsne2", "ind", "stim", "cluster", "cell", "multiplets"])
meta = meta[(meta.multiplets == "singlet") & meta.cell.notna() & (meta.cell != "NA")].copy()
meta["ind"] = meta.ind.astype(str)
M["cells"] = int(len(meta))
M["donors"] = sorted(meta.ind.unique().tolist())
M["cell_types"] = meta.cell.value_counts().to_dict()

# ---------- ① 常见错误:把所有细胞倒在一起,按细胞数做卡方 ----------
pooled = pd.crosstab(meta.cell, meta.stim)
naive = []
for ct in pooled.index:
    a, b = int(pooled.loc[ct, "ctrl"]), int(pooled.loc[ct, "stim"])
    rest_a, rest_b = int(pooled["ctrl"].sum() - a), int(pooled["stim"].sum() - b)
    chi2, p, _, _ = stats.chi2_contingency([[a, rest_a], [b, rest_b]])
    naive.append({"cell_type": ct, "ctrl_cells": a, "stim_cells": b, "p_pooled": float(p)})
naive = pd.DataFrame(naive)
naive["padj_pooled"] = np.minimum(naive.p_pooled * len(naive), 1.0)          # Bonferroni,保守
M["naive_significant"] = int((naive.padj_pooled < 0.05).sum())

# ---------- ② 正确做法:每位供者一个比例,供者层面配对检验 ----------
counts = meta.groupby(["ind", "stim", "cell"]).size().rename("n").reset_index()
tot = counts.groupby(["ind", "stim"])["n"].transform("sum")
counts["prop"] = counts.n / tot
wide = counts.pivot_table(index=["cell", "ind"], columns="stim", values="prop").fillna(0).reset_index()

proper = []
for ct, g in wide.groupby("cell"):
    if len(g) < 5:                      # 供者太少不做检验,如实标出
        proper.append({"cell_type": ct, "donors": len(g), "p_donor": None}); continue
    # 配对 Wilcoxon:同一个人自己跟自己比,把个体差异吃掉
    try:
        stat, p = stats.wilcoxon(g["ctrl"], g["stim"])
    except ValueError:
        p = 1.0
    proper.append({
        "cell_type": ct, "donors": int(len(g)),
        "mean_ctrl_pct": round(100 * float(g["ctrl"].mean()), 2),
        "mean_stim_pct": round(100 * float(g["stim"].mean()), 2),
        "p_donor": float(p),
    })
proper = pd.DataFrame(proper)
valid = proper.p_donor.notna()
proper.loc[valid, "padj_donor"] = np.minimum(proper.loc[valid, "p_donor"] * valid.sum(), 1.0)
M["donor_level_significant"] = int((proper.padj_donor < 0.05).sum())

merged = naive.merge(proper, on="cell_type")
merged["pooled_calls_it"] = merged.padj_pooled < 0.05
merged["donor_test_calls_it"] = merged.padj_donor < 0.05
M["false_positives_from_pooling"] = int((merged.pooled_calls_it & ~merged.donor_test_calls_it).sum())
M["table"] = [{
    "cell_type": r.cell_type, "ctrl_pct": r.mean_ctrl_pct, "stim_pct": r.mean_stim_pct,
    "padj_pooled": float(f"{r.padj_pooled:.2g}"), "padj_donor": None if pd.isna(r.padj_donor) else float(f"{r.padj_donor:.2g}"),
    "pooled": bool(r.pooled_calls_it), "donor": bool(r.donor_test_calls_it),
} for r in merged.itertuples()]

# ---------- 图 ----------
fig, ax = plt.subplots(figsize=(7.4, 4.4))
x = np.arange(len(merged)); w = 0.38
ax.bar(x - w/2, merged.mean_ctrl_pct, w, label="control", color="#0a6ed1")
ax.bar(x + w/2, merged.mean_stim_pct, w, label="IFN-β", color="#dc2626")
ax.set_xticks(x); ax.set_xticklabels(merged.cell_type, rotation=35, ha="right", fontsize=9)
ax.set_ylabel("mean % of cells per donor"); ax.legend()
ax.set_title("Cell-type composition, averaged per donor", fontsize=11)
plt.tight_layout(); plt.savefig(f"{OUT}/composition.png", dpi=140); plt.close()

fig, ax = plt.subplots(figsize=(6.2, 4.4))
ax.bar(["Pooled cells\n(n = %s cells)" % f"{M['cells']:,}", "Per donor\n(n = %d donors)" % len(M["donors"])],
       [M["naive_significant"], M["donor_level_significant"]], color=["#dc2626", "#0d9488"])
for i, v in enumerate([M["naive_significant"], M["donor_level_significant"]]):
    ax.text(i, v, f" {v}", ha="center", va="bottom", fontsize=12, fontweight="bold")
ax.set_ylabel("cell types called as changed (adj. p < 0.05)")
ax.set_title("Same data, two ways of counting", fontsize=11)
plt.tight_layout(); plt.savefig(f"{OUT}/method_comparison.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", "naive_significant", "donor_level_significant",
                                     "false_positives_from_pooling")}, indent=2))
for r in M["table"]:
    print(f"  {r['cell_type']:<20} ctrl {r['ctrl_pct']:>5}%  stim {r['stim_pct']:>5}%  "
          f"pooled_padj={r['padj_pooled']:<10} donor_padj={r['padj_donor']}")