#!/usr/bin/env python
"""bulk 数据的细胞组成解卷积 —— 而且带「标准答案」。

思路:用公开单细胞数据(GSE96583 对照组)造出**已知比例**的人工 bulk 样本,
再假装不知道比例去解卷积,最后和真值比。这样能给出一个别人给不了的东西:**准确度**。

为什么值得单独做一份:客户买解卷积时,拿到的通常只是一张比例饼图,
没有任何东西说明这张图准不准。而只要用留出的细胞造混合样本,准确度是可以直接量出来的。

参考集与测试集**按供者分开**(而不是随机分细胞),否则同一个人的细胞既在参考里又在测试里,
准确度会被高估 —— 这正是解卷积论文里最常见的自欺方式。
"""
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
from scipy.optimize import nnls

HERE = os.path.dirname(os.path.abspath(__file__))
D = os.path.join(HERE, "kang")
OUT = os.path.join(HERE, "out_deconv")
os.makedirs(OUT, exist_ok=True)
rng = np.random.default_rng(0)
M = {"dataset": "GEO GSE96583 (Kang et al. 2018) — control PBMC, 8 donors"}

genes = np.array([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"])
X = sio.mmread(f"{D}/ctrl.mtx.gz").tocsr()
bc = [l.strip() for l in gzip.open(f"{D}/ctrl.barcodes.tsv.gz", "rt")]
sub = meta[meta.stim == "ctrl"]
idx = [i for i, b in enumerate(bc) if b in sub.index]
X = X[:, idx].T.tocsr()
md = sub.loc[[bc[i] for i in idx]].copy()
keep = (md.multiplets == "singlet") & md.cell.notna() & (md.cell != "NA")
X, md = X[keep.values], md[keep].copy()
md["ind"] = md.ind.astype(str)

donors = sorted(md.ind.unique())
ref_donors, test_donors = donors[:4], donors[4:]
M["reference_donors"], M["test_donors"] = ref_donors, test_donors
M["cells"] = int(X.shape[0])

types = sorted(md.cell.unique())
M["cell_types"] = types

# ---------- 参考签名:只用参考供者的细胞 ----------
ref_mask = md.ind.isin(ref_donors).values
sig = np.vstack([
    np.asarray(X[ref_mask & (md.cell == t).values].mean(axis=0)).ravel() for t in types
]).T                                            # 基因 × 细胞类型
# 只留在细胞类型间有区分度的基因,签名矩阵才不会病态
expressed = sig.sum(axis=1) > 0
score = np.zeros(sig.shape[0])
score[expressed] = sig[expressed].max(axis=1) / (sig[expressed].mean(axis=1) + 1e-9)
top_genes = np.argsort(score)[::-1][:2000]
S = sig[top_genes]
S = S / (S.sum(axis=0, keepdims=True) + 1e-9)   # 每种细胞类型归一化到等总量
M["signature_genes"] = int(len(top_genes))

# ---------- 造 30 个已知比例的人工 bulk 样本(只用测试供者的细胞) ----------
test_mask = md.ind.isin(test_donors).values
Xt, mdt = X[test_mask], md[test_mask]
by_type = {t: np.flatnonzero((mdt.cell == t).values) for t in types}
N_MIX, N_CELLS = 30, 800
truths, mixes = [], []
for _ in range(N_MIX):
    w = rng.dirichlet(np.ones(len(types)) * 0.7)          # 随机但不均匀的比例
    picks = []
    for t, p in zip(types, w):
        k = int(round(p * N_CELLS))
        pool = by_type[t]
        if k > 0 and len(pool):
            picks.append(rng.choice(pool, size=min(k, len(pool)), replace=len(pool) < k))
    picks = np.concatenate(picks) if picks else np.array([], dtype=int)
    真 = np.array([np.sum(mdt.cell.values[picks] == t) for t in types], dtype=float)
    真 = 真 / 真.sum()
    truths.append(真)
    mixes.append(np.asarray(Xt[picks].sum(axis=0)).ravel()[top_genes])
truths = np.vstack(truths)
B = np.vstack(mixes).T
B = B / (B.sum(axis=0, keepdims=True) + 1e-9)
M["mixtures"] = N_MIX
M["cells_per_mixture"] = N_CELLS

# ---------- 解卷积:非负最小二乘 ----------
est = np.zeros_like(truths)
for i in range(N_MIX):
    coef, _ = nnls(S, B[:, i])
    est[i] = coef / (coef.sum() + 1e-9)

err = est - truths
M["overall_pearson_r"] = round(float(np.corrcoef(est.ravel(), truths.ravel())[0, 1]), 3)
M["overall_rmse_pct"] = round(float(np.sqrt((err ** 2).mean()) * 100), 2)
M["per_type"] = [{
    "cell_type": t,
    "mean_true_pct": round(float(truths[:, j].mean() * 100), 1),
    "mean_estimated_pct": round(float(est[:, j].mean() * 100), 1),
    "rmse_pct": round(float(np.sqrt((err[:, j] ** 2).mean()) * 100), 2),
    "r": round(float(np.corrcoef(est[:, j], truths[:, j])[0, 1]), 3) if truths[:, j].std() > 0 else None,
} for j, t in enumerate(types)]
M["worst_type"] = max(M["per_type"], key=lambda d: d["rmse_pct"])["cell_type"]

# ---------- 图 ----------
fig, ax = plt.subplots(figsize=(5.4, 5.0))
cmap = plt.get_cmap("tab10")
for j, t in enumerate(types):
    ax.scatter(truths[:, j] * 100, est[:, j] * 100, s=22, alpha=.75, color=cmap(j % 10), label=t)
lim = max(truths.max(), est.max()) * 100 * 1.05
ax.plot([0, lim], [0, lim], color="#94a3b8", lw=1, ls="--")
ax.set_xlabel("true proportion (%)"); ax.set_ylabel("estimated proportion (%)")
ax.set_title(f"Deconvolution accuracy (r = {M['overall_pearson_r']})", fontsize=11)
ax.legend(fontsize=7, frameon=False, loc="upper left")
plt.tight_layout(); plt.savefig(f"{OUT}/accuracy.png", dpi=140); plt.close()

fig, ax = plt.subplots(figsize=(6.6, 4.0))
d = pd.DataFrame(M["per_type"]).sort_values("rmse_pct")
ax.barh(d.cell_type, d.rmse_pct, color="#0a6ed1")
ax.set_xlabel("RMSE (percentage points)")
ax.set_title("Which cell types are estimated reliably, and which are not", fontsize=11)
plt.tight_layout(); plt.savefig(f"{OUT}/per_type_error.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", "reference_donors", "test_donors", "mixtures",
                                     "signature_genes", "overall_pearson_r", "overall_rmse_pct",
                                     "worst_type")}, indent=2, ensure_ascii=False))
for r in M["per_type"]:
    print(f"  {r['cell_type']:<20} 真值 {r['mean_true_pct']:>5}%  估计 {r['mean_estimated_pct']:>5}%  "
          f"RMSE {r['rmse_pct']:>5}pp  r={r['r']}")