#!/usr/bin/env python
"""CITE-seq(RNA + 表面蛋白)样例 —— 公开数据:10x 5k PBMC with TotalSeq-B。

为什么补这一份:蛋白+RNA 的项目单价高、客户多,而它有一套单独的质控要点,
拿单细胞 RNA 那套照搬会漏掉最关键的东西:

  ① **同型对照(isotype control)**:抗体非特异结合的背景。不看它,任何「阳性」都可能是背景;
  ② 蛋白与 RNA 不是一回事:蛋白半衰期长、RNA 转瞬即逝,CD4 这类基因 RNA 常测不到而蛋白很亮 ——
     所以两者不一致往往是生物学,不是错误;
  ③ 蛋白通道的标准化要用 CLR(按细胞对各抗体做中心化对数比),不能直接套 RNA 的做法。

自检:用 RNA 聚出来的 T/B/单核细胞群,在蛋白通道上 CD3/CD19/CD14 必须各自点亮对应群 ——
两条独立测量互相印证,才说明流程没问题。
"""
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 scanpy as sc

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(HERE, "out_cite")
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 5k PBMC with TotalSeq-B (public CITE-seq)"}

H5 = os.path.join(HERE, "data", "5k_pbmc_protein_v3_filtered_feature_bc_matrix.h5")
adata = sc.read_10x_h5(H5, gex_only=False)
adata.var_names_make_unique()
is_adt = (adata.var["feature_types"] == "Antibody Capture").values
rna = adata[:, ~is_adt].copy()
adt = adata[:, is_adt].copy()
M["cells"] = int(adata.n_obs)
M["rna_genes"] = int(rna.n_vars)
M["antibodies"] = [str(x) for x in adt.var_names]
M["n_antibodies"] = int(adt.n_vars)

# ---------- ① 同型对照:蛋白通道的背景水平 ----------
iso = [g for g in adt.var_names if "isotype" in str(g).lower() or "IgG" in str(g)]
M["isotype_controls"] = iso
adt_counts = np.asarray(adt.X.todense()) if hasattr(adt.X, "todense") else np.asarray(adt.X)
adt_df = pd.DataFrame(adt_counts, columns=adt.var_names)
if iso:
    iso_med = float(np.median(adt_df[iso].sum(axis=1)))
    tot_med = float(np.median(adt_df.sum(axis=1)))
    M["isotype_median_counts"] = round(iso_med, 1)
    M["adt_total_median_counts"] = round(tot_med, 1)
    M["isotype_share_pct"] = round(100 * iso_med / max(tot_med, 1), 2)

# ---------- RNA 侧的常规流程 ----------
rna.var["mt"] = rna.var_names.str.startswith("MT-")
sc.pp.calculate_qc_metrics(rna, qc_vars=["mt"], percent_top=None, log1p=False, inplace=True)
sc.pp.filter_cells(rna, min_genes=200)
rna = rna[rna.obs.pct_counts_mt < 15].copy()
M["cells_after_qc"] = int(rna.n_obs)
sc.pp.normalize_total(rna, target_sum=1e4); sc.pp.log1p(rna)
sc.pp.highly_variable_genes(rna, min_mean=0.0125, max_mean=3, min_disp=0.5)
rna.raw = rna
rna = rna[:, rna.var.highly_variable].copy()
sc.pp.scale(rna, max_value=10)
sc.tl.pca(rna, svd_solver="arpack", n_comps=50)
sc.pp.neighbors(rna, n_neighbors=10, n_pcs=30)
sc.tl.umap(rna)
sc.tl.leiden(rna, resolution=0.5, key_added="leiden", flavor="igraph", n_iterations=2, directed=False)
M["clusters"] = int(rna.obs.leiden.nunique())

# ---------- ③ 蛋白通道:CLR 标准化(不是 RNA 那套)----------
adt_sub = adt_df.loc[[adata.obs_names.get_loc(b) for b in rna.obs_names]].reset_index(drop=True)
clr = np.log1p(adt_sub.values / np.exp(np.log1p(adt_sub.values).mean(axis=1, keepdims=True)))
clr_df = pd.DataFrame(clr, columns=adt.var_names)
for c in clr_df.columns:
    rna.obs["ADT_" + str(c)] = clr_df[c].values
M["adt_normalisation"] = "CLR (centred log-ratio per cell), as recommended for antibody-derived tags"

# ---------- ② 自检:RNA 分群 vs 蛋白信号 ----------
PAIRS = [("CD3", "CD3E"), ("CD19", "MS4A1"), ("CD14", "CD14"), ("CD8a", "CD8A"), ("CD4", "IL7R")]
rows = []
for ab, gene in PAIRS:
    hit = [c for c in adt.var_names if str(c).upper().startswith(ab.upper())]
    if not hit or gene not in rna.raw.var_names: continue
    ab_col = "ADT_" + str(hit[0])
    prot = rna.obs[ab_col].values
    expr = np.asarray(rna.raw[:, gene].X.todense()).ravel()
    # 每个 RNA 群里两者的平均值,再看跨群相关 —— 群水平比单细胞水平稳
    g = pd.DataFrame({"cl": rna.obs.leiden.values, "prot": prot, "rna": expr}).groupby("cl").mean()
    r = float(np.corrcoef(g.prot, g.rna)[0, 1]) if g.prot.std() > 0 and g.rna.std() > 0 else float("nan")
    top_prot = str(g.prot.idxmax()); top_rna = str(g.rna.idxmax())
    rows.append({"antibody": str(hit[0]), "rna_gene": gene, "cluster_level_r": round(r, 3),
                 "top_cluster_protein": top_prot, "top_cluster_rna": top_rna,
                 "same_cluster": top_prot == top_rna})
M["protein_vs_rna"] = rows
M["pairs_agreeing"] = int(sum(1 for r in rows if r["same_cluster"]))
M["pairs_checked"] = len(rows)

sc.pl.umap(rna, color="leiden", legend_loc="on data", title="RNA-based clusters",
           save="_cite_clusters.png", show=False)
show = [("ADT_" + str(c)) for c in adt.var_names if str(c).upper().startswith(("CD3", "CD19", "CD14", "CD8"))][:4]
if show:
    sc.pl.umap(rna, color=show, ncols=4, save="_cite_protein.png", show=False)

fig, ax = plt.subplots(figsize=(6.6, 4.0))
tot = adt_df.sum(axis=0).sort_values(ascending=False)
colors = ["#dc2626" if c in iso else "#0a6ed1" for c in tot.index]
ax.bar(range(len(tot)), np.log10(tot.values + 1), color=colors)
ax.set_xticks(range(len(tot))); ax.set_xticklabels(tot.index, rotation=70, ha="right", fontsize=6.5)
ax.set_ylabel("log10 total counts"); ax.set_title("Antibody signal (red = isotype controls)", fontsize=11)
plt.tight_layout(); plt.savefig(f"{OUT}/adt_levels.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", "cells_after_qc", "n_antibodies", "isotype_controls",
                                     "isotype_share_pct", "clusters", "pairs_agreeing", "pairs_checked")},
                 indent=2, ensure_ascii=False))
for r in rows: print("  ", r)