#!/usr/bin/env python
"""细胞间通讯分析样例 —— 按「比较两种条件」而不是「读绝对分数」来做。

数据:GEO GSE96583(8 位供者,对照 vs IFN-β),本地已有。

我们在指南里写过:配体-受体打分只能说明「配体在 A 群表达、受体在 B 群表达」,
既不能证明两群挨着,也不能证明蛋白被翻译和结合。所以能站得住的用法是**比较条件**,
而不是把绝对分数当成「这条通路是活跃的」。这份样品就按那个用法跑,并给出自检:
干扰素刺激后,以 CXCL10/CXCL11 等干扰素诱导趋化因子为配体的互作应当明显增强。
"""
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 scipy.io as sio
import scipy.sparse as sp
import gzip
import anndata as ad
import scanpy as sc

HERE = os.path.dirname(os.path.abspath(__file__))
D = os.path.join(HERE, "kang")
OUT = os.path.join(HERE, "out_comm")
os.makedirs(OUT, exist_ok=True)
M = {"dataset": "GEO GSE96583 (Kang et al. 2018) — PBMC, 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"])

def load(tag, mtx, bcs):
    X = sio.mmread(f"{D}/{mtx}").tocsr()
    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]
    m = sub.loc[[bc[i] for i in idx]].copy(); m["condition"] = tag
    return X[:, idx].T.tocsr(), 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])
keep = (md.multiplets == "singlet") & md.cell.notna() & (md.cell != "NA")
X, md = X[keep.values], md[keep]

A = ad.AnnData(X=X.astype(np.float32), obs=md.reset_index()[["cell", "condition", "ind"]].astype(str))
A.var_names = pd.Index(genes); A.var_names_make_unique()
A.obs.rename(columns={"cell": "cell_type"}, inplace=True)
sc.pp.filter_genes(A, min_cells=20)
sc.pp.normalize_total(A, target_sum=1e4); sc.pp.log1p(A)
M["cells"] = int(A.n_obs); M["cell_types"] = sorted(A.obs.cell_type.unique().tolist())

import liana as li
res = {}
for cond in ["ctrl", "stim"]:
    sub = A[A.obs.condition == cond].copy()
    li.mt.rank_aggregate(sub, groupby="cell_type", expr_prop=0.1, use_raw=False, verbose=False)
    r = sub.uns["liana_res"].copy()
    r["pair"] = r.source + " → " + r.target + " | " + r.ligand_complex + "-" + r.receptor_complex
    res[cond] = r
    M[f"{cond}_interactions_tested"] = int(len(r))

# 「显著」按 liana 的聚合秩:两边都取排名最靠前的一批,再比差异
def top(r, n=200):
    return r.nsmallest(n, "magnitude_rank")[["pair", "magnitude_rank", "specificity_rank",
                                             "source", "target", "ligand_complex", "receptor_complex"]]
tc, ts = top(res["ctrl"]), top(res["stim"])
only_stim = ts[~ts.pair.isin(tc.pair)]
only_ctrl = tc[~tc.pair.isin(ts.pair)]
M["top_n_compared"] = 200
M["gained_with_ifn"] = int(len(only_stim))
M["lost_with_ifn"] = int(len(only_ctrl))
M["shared"] = int(200 - len(only_stim))

ISG_LIG = ["CXCL10", "CXCL11", "CXCL9", "IL15", "TNFSF10", "ICAM1", "HLA-A", "HLA-B", "HLA-C", "HLA-E", "B2M"]
gained_isg = only_stim[only_stim.ligand_complex.isin(ISG_LIG)]
M["ifn_induced_ligands_gained"] = sorted(gained_isg.ligand_complex.unique().tolist())
M["sanity_check_passed"] = bool(len(M["ifn_induced_ligands_gained"]) > 0)
M["top_gained"] = [{"pair": r.pair, "rank": float(f"{r.magnitude_rank:.3g}")}
                   for r in only_stim.head(12).itertuples()]

# 图:各细胞类型在两种条件下进入 top200 的互作数
cnt = pd.DataFrame({
    "ctrl": tc.source.value_counts(),
    "stim": ts.source.value_counts(),
}).fillna(0)
cnt = cnt.loc[cnt.sum(axis=1).sort_values(ascending=False).index]
fig, ax = plt.subplots(figsize=(7.2, 4.2))
x = np.arange(len(cnt)); w = .38
ax.bar(x - w/2, cnt.ctrl, w, label="control", color="#0a6ed1")
ax.bar(x + w/2, cnt.stim, w, label="IFN-β", color="#dc2626")
ax.set_xticks(x); ax.set_xticklabels(cnt.index, rotation=30, ha="right", fontsize=9)
ax.set_ylabel("interactions in top 200 (as sender)"); ax.legend()
ax.set_title("Signalling output by cell type, compared between conditions", fontsize=11)
plt.tight_layout(); plt.savefig(f"{OUT}/by_celltype.png", dpi=140); plt.close()

fig, ax = plt.subplots(figsize=(6.6, 4.4))
g = gained_isg.ligand_complex.value_counts().head(10)[::-1]
if len(g):
    ax.barh(g.index, g.values, color="#dc2626")
    ax.set_xlabel("interactions gained after IFN-β (top 200)")
    ax.set_title("Interferon-induced ligands appearing after stimulation", fontsize=11)
    plt.tight_layout(); plt.savefig(f"{OUT}/gained_ligands.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", "gained_with_ifn", "lost_with_ifn", "shared",
                                     "ifn_induced_ligands_gained", "sanity_check_passed")},
                 indent=2, ensure_ascii=False))
for t in M["top_gained"][:8]:
    print("  ", t["pair"])