#!/usr/bin/env python
"""自动注释 × 人工 marker 注释的一致性检验(公开 PBMC 3k 数据)。

为什么做这个:注释是单细胞分析里最容易被审稿人挑战的一步,而我们在指南里主张
「三条独立证据汇合才下结论」。样品页上光说没用 —— 这个脚本真跑一遍:
用 CellTypist 的公开免疫模型(完全不看我们的人工标签)独立注释,再和 marker 法的结果对表。
一致就是证据,不一致的地方如实列出来,那才是真正该讨论的细胞群。
"""
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

OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "out_annot")
os.makedirs(OUT, exist_ok=True)
sc.settings.figdir = OUT
sc.settings.set_figure_params(dpi=140, frameon=False, figsize=(5.4, 4.4), facecolor="white")
M = {}

# ---------- 复现样品页那套人工流程 ----------
adata = sc.datasets.pbmc3k()
adata.var["mt"] = adata.var_names.str.startswith("MT-")
sc.pp.calculate_qc_metrics(adata, qc_vars=["mt"], percent_top=None, log1p=False, inplace=True)
sc.pp.filter_cells(adata, min_genes=200)
sc.pp.filter_genes(adata, min_cells=3)
adata = adata[(adata.obs.n_genes_by_counts < 2500) & (adata.obs.pct_counts_mt < 5), :].copy()
sc.pp.normalize_total(adata, target_sum=1e4)
sc.pp.log1p(adata)
lognorm = adata.copy()                      # CellTypist 要的就是这个尺度(1e4 + log1p)

sc.pp.highly_variable_genes(adata, min_mean=0.0125, max_mean=3, min_disp=0.5)
adata.raw = adata
adata = adata[:, adata.var.highly_variable].copy()
sc.pp.regress_out(adata, ["total_counts", "pct_counts_mt"])
sc.pp.scale(adata, max_value=10)
sc.tl.pca(adata, svd_solver="arpack", n_comps=50)
sc.pp.neighbors(adata, n_neighbors=10, n_pcs=30)
sc.tl.umap(adata)
sc.tl.leiden(adata, resolution=0.5, key_added="leiden", flavor="igraph", n_iterations=2, directed=False)
sc.tl.rank_genes_groups(adata, "leiden", method="wilcoxon")
top = {g: [str(x) for x in adata.uns["rank_genes_groups"]["names"][g][:8]]
       for g in adata.obs["leiden"].cat.categories}

SIG = {
    "CD4+ T": ["IL7R", "CCR7", "CD3D"], "CD8+ T": ["CD8A", "CD8B", "GZMK"],
    "NK": ["GNLY", "NKG7", "KLRD1"], "B": ["MS4A1", "CD79A", "CD79B"],
    "CD14+ Monocyte": ["CD14", "LYZ", "S100A9"], "FCGR3A+ Monocyte": ["FCGR3A", "MS4A7"],
    "Dendritic": ["FCER1A", "CST3"], "Platelet": ["PPBP", "PF4"],
}
manual = {}
for cl, marks in top.items():
    best, hit = "Unassigned", 0
    for name, sig in SIG.items():
        k = len(set(sig) & set(marks))
        if k > hit: best, hit = name, k
    manual[cl] = best if hit else "Unassigned"
adata.obs["manual"] = adata.obs["leiden"].map(manual).astype("category")

# ---------- 自动注释:完全不看人工标签 ----------
import celltypist
from celltypist import models
models.download_models(force_update=False, model=["Immune_All_Low.pkl"])
pred = celltypist.annotate(lognorm[adata.obs_names].copy(), model="Immune_All_Low.pkl",
                           majority_voting=True, over_clustering=adata.obs["leiden"].astype(str).values)
auto = pred.predicted_labels["majority_voting"].astype(str)
adata.obs["auto"] = pd.Categorical(auto.values)
M["model"] = "CellTypist Immune_All_Low (public)"

# ---------- 对表 ----------
rows = []
for cl in adata.obs["leiden"].cat.categories:
    m = adata.obs["leiden"] == cl
    lab = adata.obs.loc[m, "auto"].value_counts()
    rows.append({
        "cluster": str(cl), "cells": int(m.sum()),
        "manual": manual[cl],
        "automated": str(lab.index[0]),
        "automated_agreement": round(100 * float(lab.iloc[0]) / int(m.sum()), 1),
        "top_markers": ", ".join(top[cl][:5]),
    })
M["clusters"] = rows

# 粗粒度谱系对齐:自动标签用的是更细的免疫细胞名,所以按谱系归并再判一致
def lineage(s):
    # ⚠️ 谱系归并要能吃下两边的写法:人工标签很短("B"),自动标签很长("Tcm/Naive helper T cells")。
    # 第一版漏了单字母 "B",把本来一致的 B 细胞群判成不一致 —— 是归并函数的 bug,不是注释不一致。
    s = s.lower().strip()
    if s in ("b", "b cell", "b cells"): return "B/plasma"
    if s in ("nk",): return "NK"
    if "cd8" in s or "cytotoxic" in s or "mait" in s: return "CD8/cytotoxic T"
    if "treg" in s or "cd4" in s or "helper" in s or " t cell" in s or s.startswith("t cell") or "tcm" in s or "tem" in s: return "CD4/other T"
    if "nk" in s: return "NK"
    if "b cell" in s or s.startswith("b ") or "plasma" in s or "memory b" in s or "naive b" in s: return "B/plasma"
    if "monocyt" in s or "macrophage" in s: return "Monocyte/macrophage"
    if "dc" in s or "dendritic" in s: return "Dendritic"
    if "platelet" in s or "megakaryo" in s: return "Platelet"
    return "other"

agree = sum(1 for r in rows if lineage(r["manual"]) == lineage(r["automated"]))
M["clusters_total"] = len(rows)
M["clusters_lineage_agree"] = agree
M["agreement_pct"] = round(100 * agree / len(rows), 1)
M["disagreements"] = [r for r in rows if lineage(r["manual"]) != lineage(r["automated"])]

sc.pl.umap(adata, color="manual", title="Manual, marker-based", save="_manual.png", show=False)
sc.pl.umap(adata, color="auto", title="Automated (CellTypist, blind)", save="_auto.png", show=False)

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 ("model", "clusters_total", "clusters_lineage_agree", "agreement_pct")}, indent=2))
for r in rows:
    print(f"  cluster {r['cluster']:>2} n={r['cells']:>5}  manual={r['manual']:<18} auto={r['automated']:<28} ({r['automated_agreement']}%)")