# -*- coding: utf-8 -*-
"""gen_figures_partB_analysis.py —— 生成"模拟数据分析结果"风格教学图 fig07–fig17。

思路：先在代码内生成一个"模拟的 10x Visium 人乳腺癌样组织"数据集
（六边形 spot 网格 + 四个空间生态位 + 六种细胞类型的空间组成 + 表达矩阵），
再用标准空间转录组分析流程（QC → PCA/聚类 → 注释 → 生态位 → 邻域富集 → CellChat
风格通讯 → 差异通讯/重编程）绘制出与真实分析高度一致的论文风格图。
所有图注均声明"模拟数据演示"。

运行：python code/figures/gen_figures_partB_analysis.py
输出：figures/fig07_qc.png ... figures/fig17_paper_figure.png
      data/sim_visium.npz（模拟数据集，供复现）
"""
import os, sys
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrowPatch, Circle, FancyBboxPatch
from matplotlib.path import Path
import matplotlib.patches as mpatches
from matplotlib.colors import LinearSegmentedColormap, Normalize
from matplotlib import cm
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from fig_common import savefig, NICHE_COLORS, CT_COLORS, PALETTE10

RNG = np.random.default_rng(42)
DATA_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "data")

# ============================================================ 1. 模拟数据集
NICHES = ["肿瘤核心", "侵袭前沿", "免疫浸润区", "基质区"]
CTS = ["肿瘤细胞", "CAF", "CD8 T", "巨噬细胞", "B 细胞", "内皮细胞"]

# 各生态位细胞组成（行=生态位，列=细胞类型），与 fig06 一致
NICHE_COMP = np.array([
    [0.62, 0.10, 0.05, 0.08, 0.03, 0.12],   # 肿瘤核心
    [0.40, 0.22, 0.12, 0.12, 0.06, 0.08],   # 侵袭前沿
    [0.18, 0.12, 0.30, 0.20, 0.15, 0.05],   # 免疫浸润区
    [0.10, 0.35, 0.10, 0.12, 0.08, 0.25],   # 基质区
])

# 基因：每类细胞的 marker + 通讯通路基因 + 噪声基因
GENES = (["EPCAM", "KRT8", "KRT19", "CD274", "TGFB1", "VEGFA", "MIF",   # 肿瘤细胞相关
          "COL1A1", "COL3A1", "DCN", "PDGFRA", "CXCL12", "FN1",          # CAF 相关
          "CD3D", "CD8A", "GZMB", "CCL5",                                # CD8 T
          "CD68", "CD163", "C1QA", "IL10", "SPP1", "CXCL13",             # 巨噬细胞/B
          "MS4A1", "CD79A", "LGALS9",                                     # B 细胞
          "PECAM1", "VWF", "KDR", "CXCR4", "TGFBR2", "PDCD1", "CD74",
          "HAVCR2", "ITGA5", "CCR5"])
NOISE_GENES = ["GAPDH", "ACTB", "B2M", "MALAT1", "TMSB4X", "RPLP0", "EEF1A1", "HSP90AB1"]
MT_GENES = ["MT-ATP6", "MT-CO1", "MT-CYB", "MT-ND4"]

# 每个基因在各细胞类型的平均表达（行=基因，列=细胞类型）
def _mk_means():
    n_g = len(GENES); means = np.full((n_g, 6), 4.0)
    for g, t_idx, val in [
        ("EPCAM", 0, 90), ("KRT8", 0, 70), ("KRT19", 0, 60), ("CD274", 0, 22),
        ("TGFB1", 0, 25), ("TGFB1", 1, 20), ("VEGFA", 0, 30), ("MIF", 0, 45),
        ("MIF", 3, 30),
        ("COL1A1", 1, 95), ("COL3A1", 1, 80), ("DCN", 1, 55), ("PDGFRA", 1, 40),
        ("CXCL12", 1, 65), ("FN1", 1, 35), ("FN1", 3, 22),
        ("CD3D", 2, 70), ("CD8A", 2, 65), ("GZMB", 2, 45), ("CCL5", 2, 50),
        ("CD68", 3, 75), ("CD163", 3, 60), ("C1QA", 3, 50), ("IL10", 3, 28),
        ("SPP1", 3, 40), ("SPP1", 0, 15), ("CXCL13", 3, 35), ("CXCL13", 4, 25),
        ("MS4A1", 4, 75), ("CD79A", 4, 65), ("LGALS9", 4, 30), ("LGALS9", 0, 20),
        ("PECAM1", 5, 85), ("VWF", 5, 70), ("KDR", 5, 45),
        ("CXCR4", 0, 35), ("TGFBR2", 1, 30), ("PDCD1", 2, 28),
        ("CD74", 3, 35), ("HAVCR2", 2, 22), ("ITGA5", 0, 18), ("CCR5", 3, 20),
    ]:
        means[GENES.index(g), t_idx] = val
    return means

def sim_visium():
    """生成模拟 Visium 数据，返回 dict。"""
    means = _mk_means()

    # --- 六边形 spot 网格 + 组织掩膜（超椭圆，带两个"叶瓣"模拟乳腺组织） ---
    dx, dy = 0.24, 0.24 * np.sqrt(3) / 2
    spots = []
    i = 0
    for ii in range(-28, 29):
        for jj in range(-30, 31):
            x = ii * dx + (jj % 2) * dx / 2
            y = jj * dy
            # 主团块
            d1 = ((abs(x) / 6.4) ** 2.4 + (abs(y) / 5.2) ** 2.4)
            # 两个叶瓣
            d2 = ((x - 3.4) / 1.7) ** 2 + ((y + 2.6) / 1.9) ** 2
            d3 = ((x + 4.2) / 1.6) ** 2 + ((y - 3.4) / 1.8) ** 2
            if d1 < 1.0 or d2 < 1.0 or d3 < 1.0:
                spots.append((x, y, i)); i += 1
    spots = np.array(spots)
    n = len(spots)
    X, Y = spots[:, 0], spots[:, 1]

    # --- 生态位划分 ---
    D = np.sqrt(X ** 2 + Y ** 2)
    niche = np.full(n, 3)          # 默认基质区
    niche[D < 1.7] = 0             # 肿瘤核心
    niche[(D >= 1.7) & (D < 3.0)] = 1   # 侵袭前沿
    # 免疫浸润补丁（椭圆）
    patches = [(1.2, 1.4, 1.3, 0.95, 0.5), (-2.2, 0.6, 1.4, 1.0, -0.4),
               (3.8, -1.6, 1.1, 0.85, 0.9), (-3.6, -2.8, 1.2, 0.9, -0.9),
               (0.6, -3.4, 1.0, 0.8, 0.3), (-0.8, 2.6, 1.0, 0.8, -0.3),
               (4.6, 1.6, 0.9, 0.7, 0.1), (-4.6, -0.2, 0.9, 0.7, 0.6),
               (2.6, 3.6, 0.9, 0.7, -0.6), (-2.6, 4.0, 0.8, 0.7, 0.2)]
    for (px, py, a, b, rot) in patches:
        c, s = np.cos(rot), np.sin(rot)
        dxp = (X - px) * c - (Y - py) * s
        dyp = (X - px) * s + (Y - py) * c
        m = ((dxp / a) ** 2 + (dyp / b) ** 2 < 1) & (niche != 0)
        niche[m] = 2
    # 保证每类数量（如有太少则并入邻近）
    for k in range(4):
        if (niche == k).sum() < 60:
            niche[(niche == k)] = 1

    # --- 每 spot 细胞类型比例（Dirichlet 围绕生态位组成） ---
    alpha = 55.0
    prop = np.zeros((n, 6))
    for k in range(4):
        m = niche == k
        prop[m] = RNG.dirichlet(alpha * NICHE_COMP[k] + 1e-3, size=m.sum())

    # --- 表达矩阵 ---
    n_g = len(GENES)
    spot_mean = prop @ means.T          # n × n_g
    sf = RNG.lognormal(0, 0.25, n)      # 文库大小因子
    counts = RNG.poisson(spot_mean * sf[:, None] * 1.6)
    nCounts = counts.sum(axis=1)
    nGenes = (counts > 0).sum(axis=1)

    # 线粒体比例（按生态位略不同 + 噪声）
    mito_base = np.array([0.10, 0.12, 0.08, 0.09])
    pct_mito = np.clip(mito_base[niche] + RNG.normal(0, 0.015, n), 0.02, 0.30)
    mito_counts = (pct_mito * nCounts)[:, None] * np.array([0.31, 0.27, 0.22, 0.20])
    counts = np.hstack([counts, mito_counts.astype(int)])

    # 5% 低质量 spot（QC 过滤演示）
    lowq = RNG.random(n) < 0.05
    counts[lowq] = (counts[lowq] * 0.12).astype(int)
    nCounts = counts.sum(axis=1)
    nGenes = (counts > 0).sum(axis=1)
    pct_mito = np.clip(pct_mito + (lowq * 0.08), 0, 0.45)

    # --- 归一化（CPM → log1p） ---
    cpm = counts / nCounts[:, None] * 1e4
    logcpm = np.log1p(cpm)

    # --- PCA（基于 marker 基因，去掉噪声） ---
    Xm = logcpm[:, :n_g]
    Xs = (Xm - Xm.mean(0)) / (Xm.std(0) + 1e-9)
    U, S, Vt = np.linalg.svd(Xs, full_matrices=False)
    pca = U[:, :10] * S[:10]

    # --- 聚类（KMeans，k=8，模拟 Louvain/Leiden 结果） ---
    from sklearn.cluster import KMeans
    km = KMeans(n_clusters=8, random_state=0, n_init=10).fit(pca)
    cluster = km.labels_
    # 按与生态位的对应重新排序编号（美观）
    order = np.zeros(8, dtype=int)
    for kk in range(8):
        order[kk] = int(np.bincount(niche[cluster == kk], minlength=4).argmax())
    remap = np.zeros(8, dtype=int)
    for new, old in enumerate(np.argsort(order)):
        remap[old] = new
    cluster = remap[cluster]

    return dict(X=X, Y=Y, niche=niche, prop=prop, counts=counts,
                nCounts=nCounts, nGenes=nGenes, pct_mito=pct_mito,
                logcpm=logcpm, pca=pca, cluster=cluster, lowq=lowq,
                genes=GENES + MT_GENES, niche_names=NICHES, cts=CTS)

# ============================================================ 2. 通讯模拟
# 配体-受体对：sender, ligand, receiver, receptor, pathway, base_prob
LR_PAIRS = [
    ("肿瘤细胞", "TGFB1", "CAF", "TGFBR2", "TGFb", 0.55),
    ("肿瘤细胞", "CD274", "CD8 T", "PDCD1", "PD-L1", 0.45),
    ("肿瘤细胞", "MIF", "巨噬细胞", "CD74", "MIF", 0.50),
    ("肿瘤细胞", "LGALS9", "CD8 T", "HAVCR2", "GALECTIN", 0.38),
    ("肿瘤细胞", "VEGFA", "内皮细胞", "KDR", "VEGF", 0.42),
    ("CAF", "CXCL12", "肿瘤细胞", "CXCR4", "CXCL", 0.48),
    ("CAF", "FN1", "肿瘤细胞", "ITGA5", "ECM", 0.35),
    ("巨噬细胞", "IL10", "CD8 T", "IL10RA", "IL10", 0.30),
    ("CD8 T", "CCL5", "巨噬细胞", "CCR5", "CCL", 0.40),
    ("巨噬细胞", "SPP1", "肿瘤细胞", "CD44", "ECM", 0.32),
    ("巨噬细胞", "CXCL13", "B 细胞", "CXCR5", "CXCL", 0.28),
]
LR_GENES = set(lr[1] for lr in LR_PAIRS) | set(lr[3] for lr in LR_PAIRS)

# 各生态位的通路活性乘子
NICHE_PATHWAY_ACT = {
    "肿瘤核心": {"TGFb": 0.8, "PD-L1": 0.5, "MIF": 0.7, "GALECTIN": 0.4, "VEGF": 0.9,
                "CXCL": 0.7, "ECM": 0.6, "IL10": 0.4, "CCL": 0.3},
    "侵袭前沿": {"TGFb": 0.95, "PD-L1": 0.9, "MIF": 0.9, "GALECTIN": 0.85, "VEGF": 0.7,
                "CXCL": 0.85, "ECM": 0.85, "IL10": 0.7, "CCL": 0.5},
    "免疫浸润区": {"TGFb": 0.4, "PD-L1": 0.65, "MIF": 0.8, "GALECTIN": 0.6, "VEGF": 0.3,
                 "CXCL": 0.5, "ECM": 0.3, "IL10": 0.5, "CCL": 0.9},
    "基质区": {"TGFb": 0.5, "PD-L1": 0.3, "MIF": 0.5, "GALECTIN": 0.4, "VEGF": 0.5,
              "CXCL": 0.6, "ECM": 0.7, "IL10": 0.3, "CCL": 0.4},
}

def comm_prob_matrix(niche_idx):
    """返回 (6×6) 通讯概率矩阵 与 (LR 对 × prob) 明细，针对某生态位。"""
    act = NICHE_PATHWAY_ACT[NICHES[niche_idx]]
    M = np.zeros((6, 6))
    detail = []
    for (s, lig, r, rec, pw, bp) in LR_PAIRS:
        p = bp * act[pw]
        M[CTS.index(s), CTS.index(r)] += p
        detail.append((s, lig, r, rec, pw, p))
    return M, detail

# ============================================================ 3. 各图
def fig07_qc(data):
    fig, axes = plt.subplots(2, 3, figsize=(13.5, 7.6))
    niche = data["niche"]; colors = [NICHE_COLORS[NICHES[k]] for k in niche]
    pos_ok = ~data["lowq"]

    # 第一行：violin（按生态位）
    import matplotlib as mpl
    for ax, val, name in [(axes[0, 0], data["nCounts"], "每个 spot 的 UMI 总数"),
                          (axes[0, 1], data["nGenes"], "每个 spot 检测到的基因数"),
                          (axes[0, 2], data["pct_mito"] * 100, "线粒体基因比例 (%)")]:
        parts = ax.violinplot([val[niche == k] for k in range(4)],
                              positions=[0, 1, 2, 3], showmedians=True, widths=0.75)
        for i, b in enumerate(parts["bodies"]):
            b.set_facecolor(list(NICHE_COLORS.values())[i]); b.set_alpha(0.65)
        ax.set_xticks(range(4)); ax.set_xticklabels(NICHES, fontsize=9, rotation=12)
        ax.set_title(name, fontsize=11.5)
        ax.set_ylabel("数值", fontsize=9.5)
        ax.grid(axis="y", alpha=0.25)

    # 第二行：空间散点
    for ax, val, name, cmap in [(axes[1, 0], data["nCounts"], "空间分布：UMI 总数", "YlOrRd"),
                                (axes[1, 1], data["nGenes"], "空间分布：基因数", "YlGnBu"),
                                (axes[1, 2], data["pct_mito"] * 100, "空间分布：线粒体比例", "PuRd")]:
        sc = ax.scatter(data["X"][pos_ok], data["Y"][pos_ok], c=val[pos_ok], s=8,
                        cmap=cmap, edgecolors="none")
        ax.scatter(data["X"][~pos_ok], data["Y"][~pos_ok], s=8, c="#BDC3C7", edgecolors="none",
                   label="过滤掉的低质量 spot")
        cb = fig.colorbar(sc, ax=ax, shrink=0.8); cb.ax.tick_params(labelsize=8)
        ax.set_title(name, fontsize=11.5); ax.set_aspect("equal")
        ax.set_xticks([]); ax.set_yticks([])
        if ax is axes[1, 2]:
            ax.legend(fontsize=8, loc="lower left")

    fig.suptitle("模拟数据演示：spot 质控（QC）指标——按生态位分布与空间分布", fontsize=14, fontweight="bold", y=0.99)
    fig.tight_layout(rect=[0, 0, 1, 0.97])
    return savefig(fig, "fig07_qc.png")

def fig08_clustering(data):
    fig, axes = plt.subplots(1, 2, figsize=(12.5, 5.2))
    clust_c = [PALETTE10[c % 10] for c in data["cluster"]]

    ax = axes[0]
    ax.scatter(data["pca"][:, 0], data["pca"][:, 1], c=clust_c, s=9, alpha=0.85, edgecolors="none")
    ax.set_xlabel("PCA1"); ax.set_ylabel("PCA2")
    ax.set_title("降维可视化（模拟数据演示，PCA 投影示意）", fontsize=11.5)
    ax.grid(alpha=0.2)

    ax = axes[1]
    ax.scatter(data["X"], data["Y"], c=clust_c, s=9, edgecolors="none")
    ax.set_title("聚类结果映射回组织切片", fontsize=11.5)
    ax.set_aspect("equal"); ax.set_xticks([]); ax.set_yticks([])

    handles = [mpatches.Patch(color=PALETTE10[i % 10], label=f"Cluster {i+1}") for i in range(8)]
    fig.legend(handles=handles, loc="lower center", ncol=8, fontsize=9, frameon=False)
    fig.suptitle("模拟数据演示：PCA/聚类——无监督发现组织内的表达程序（细胞群）", fontsize=14, fontweight="bold", y=1.0)
    fig.tight_layout(rect=[0, 0.04, 1, 0.94])
    return savefig(fig, "fig08_clustering.png")

def fig09_celltype(data):
    dom = data["prop"].argmax(axis=1)
    dom_c = [CT_COLORS[CTS[d]] for d in dom]
    fig, axes = plt.subplots(1, 2, figsize=(12.5, 5.2))

    ax = axes[0]
    ax.scatter(data["X"], data["Y"], c=dom_c, s=9, edgecolors="none")
    ax.set_title("每个 spot 的主要细胞类型（注释结果）", fontsize=11.5)
    ax.set_aspect("equal"); ax.set_xticks([]); ax.set_yticks([])

    ax = axes[1]
    props_by_niche = np.zeros((4, 6))
    for k in range(4):
        props_by_niche[k] = data["prop"][data["niche"] == k].mean(0)
    bottom = np.zeros(4)
    for j in range(6):
        ax.bar(range(4), props_by_niche[:, j], bottom=bottom, color=CT_COLORS[CTS[j]],
               width=0.62, label=CTS[j], edgecolor="white", lw=0.5)
        bottom += props_by_niche[:, j]
    ax.set_xticks(range(4)); ax.set_xticklabels(NICHES, fontsize=9.5)
    ax.set_ylabel("细胞类型占比"); ax.set_title("各生态位的平均细胞组成（去卷积结果）", fontsize=11.5)
    ax.legend(fontsize=8, ncol=2, loc="upper right")

    handles = [mpatches.Patch(color=CT_COLORS[c], label=c) for c in CTS]
    fig.legend(handles=handles, loc="lower center", ncol=6, fontsize=9, frameon=False)
    fig.suptitle("模拟数据演示：细胞类型注释与空间分布", fontsize=14, fontweight="bold", y=1.0)
    fig.tight_layout(rect=[0, 0.05, 1, 0.94])
    return savefig(fig, "fig09_celltype.png")

def fig10_niche(data):
    fig, ax = plt.subplots(figsize=(8.2, 7.2))
    cols = [NICHE_COLORS[NICHES[k]] for k in data["niche"]]
    ax.scatter(data["X"], data["Y"], c=cols, s=11, edgecolors="none")
    ax.set_title("模拟数据演示：空间生态位识别结果", fontsize=13, fontweight="bold")
    ax.set_aspect("equal"); ax.set_xticks([]); ax.set_yticks([])
    handles = [mpatches.Patch(color=NICHE_COLORS[k], label=f"{k}（{n:.0%}）")
               for k, n in zip(NICHES, np.bincount(data["niche"], minlength=4) / len(data["niche"]))]
    ax.legend(handles=handles, loc="lower left", fontsize=10, framealpha=0.95)
    return savefig(fig, "fig10_niche.png")

def _nhood_enrichment(data, k=6):
    """kNN 邻域共现富集 z 分数（模拟 Squidpy nhood_enrichment 输出）。"""
    from scipy.spatial import cKDTree
    XY = np.column_stack([data["X"], data["Y"]])
    tree = cKDTree(XY)
    _, idx = tree.query(XY, k=k + 1)
    neigh = idx[:, 1:]
    dom = data["prop"].argmax(axis=1)
    n_ct = 6
    obs = np.zeros((n_ct, n_ct))
    for i in range(len(dom)):
        for j in neigh[i]:
            obs[dom[i], dom[j]] += 1
    obs = obs + obs.T
    total = obs.sum()
    freq = obs.sum(axis=1) / total
    exp = np.outer(freq, freq) * total
    z = np.where(exp > 0, (obs - exp) / np.sqrt(exp + 1e-9), 0)
    np.fill_diagonal(z, 0)
    return z

def fig11_nhood(data):
    z = _nhood_enrichment(data)
    fig, ax = plt.subplots(figsize=(8.4, 7.0))
    vmax = np.abs(z).max() * 0.8
    im = ax.imshow(z, cmap="RdBu_r", vmin=-vmax, vmax=vmax)
    ax.set_xticks(range(6)); ax.set_xticklabels(CTS, rotation=35, ha="right", fontsize=10)
    ax.set_yticks(range(6)); ax.set_yticklabels(CTS, fontsize=10)
    for i in range(6):
        for j in range(6):
            ax.text(j, i, f"{z[i, j]:.1f}", ha="center", va="center", fontsize=8.5,
                    color="white" if abs(z[i, j]) > vmax * 0.55 else "#2C3E50")
    ax.set_title("模拟数据演示：细胞类型邻域共现富集（z 分数）", fontsize=12.5, fontweight="bold")
    cb = fig.colorbar(im, ax=ax, shrink=0.85); cb.set_label("z 分数（正=共定位富集）", fontsize=10)
    return savefig(fig, "fig11_nhood.png")

def _circular_comm_plot(ax, M, title, edge_labels=None, max_label=5):
    """圆形通讯网络图：节点=细胞类型，边=通讯概率（CellChat netVisual_circle 风格）。"""
    n = M.shape[0]
    ang = np.linspace(0, 2 * np.pi, n, endpoint=False) - np.pi / 2
    pos = {i: (np.cos(a), np.sin(a)) for i, a in enumerate(ang)}
    ax.set_xlim(-1.55, 1.55); ax.set_ylim(-1.55, 1.55); ax.set_aspect("equal"); ax.axis("off")
    # 边（由强到弱）
    edges = [(i, j, M[i, j]) for i in range(n) for j in range(n) if M[i, j] > 0.02]
    edges.sort(key=lambda e: -e[2])
    vmax = max([e[2] for e in edges], default=1)
    for i, j, w in edges:
        x0, y0 = pos[i]; x1, y1 = pos[j]
        rad = 0.42
        if i == j:
            continue
        alpha = 0.25 + 0.7 * (w / vmax)
        lw = 0.8 + 4.5 * (w / vmax)
        arrow = FancyArrowPatch((x0 * 0.88, y0 * 0.88), (x1 * 0.88, y1 * 0.88),
                                connectionstyle=f"arc3,rad={rad}",
                                arrowstyle="-|>", mutation_scale=13,
                                color="#E74C3C" if i == 0 else "#5D6D7E",
                                lw=lw, alpha=alpha, shrinkA=0, shrinkB=0)
        ax.add_patch(arrow)
    # 节点
    for i in range(n):
        x, y = pos[i]
        ax.add_patch(Circle((x, y), 0.30, facecolor=CT_COLORS[CTS[i]], edgecolor="white", lw=2, zorder=5))
        ax.text(x * 1.42, y * 1.42, CTS[i], ha="center", va="center", fontsize=10.5,
                fontweight="bold", zorder=6)
    # 最强边标注配体
    if edge_labels:
        shown = 0
        for i, j, w in edges:
            if shown >= max_label:
                break
            x, y = (pos[i][0] + pos[j][0]) / 2, (pos[i][1] + pos[j][1]) / 2
            x *= 1.12; y *= 1.12
            ax.text(x, y, edge_labels[(i, j)], fontsize=8.5, color="#7B241C",
                    ha="center", va="center",
                    bbox=dict(boxstyle="round,pad=0.15", fc="white", ec="#E67E22", lw=0.7, alpha=0.9))
            shown += 1
    ax.set_title(title, fontsize=12, fontweight="bold")

def fig12_comm_network(data):
    # 三个生态位并排：侵袭前沿（重编程热点）、免疫浸润区、肿瘤核心
    fig, axes = plt.subplots(1, 3, figsize=(15.5, 5.6))
    niche_sel = [1, 2, 0]
    for ax, k in zip(axes, niche_sel):
        M, detail = comm_prob_matrix(k)
        labels = {(CTS.index(s), CTS.index(r)): lig for (s, lig, r, rec, pw, p) in detail}
        _circular_comm_plot(ax, M, f"{NICHES[k]} 生态位的通讯网络", edge_labels=labels)
    fig.suptitle("模拟数据演示：生态位内细胞通讯网络（CellChat netVisual_circle 风格）", fontsize=14, fontweight="bold", y=1.0)
    fig.tight_layout(rect=[0, 0, 1, 0.94])
    return savefig(fig, "fig12_comm_network.png")

def fig13_comm_barplot(data):
    n_int = []; strength = []
    for k in range(4):
        M, _ = comm_prob_matrix(k)
        n_int.append((M > 0.05).sum())
        strength.append(M.sum())
    fig, axes = plt.subplots(1, 2, figsize=(11.5, 4.6))
    bars = axes[0].bar(NICHES, n_int, color=list(NICHE_COLORS.values()), alpha=0.9, edgecolor="white")
    for b, v in zip(bars, n_int):
        axes[0].text(b.get_x() + b.get_width() / 2, v + 0.3, str(v), ha="center", fontsize=10)
    axes[0].set_ylabel("显著通讯对数（数量）")
    axes[0].set_title("各生态位：配体-受体通讯数量", fontsize=11.5)
    axes[0].grid(axis="y", alpha=0.25)

    bars = axes[1].bar(NICHES, strength, color=list(NICHE_COLORS.values()), alpha=0.9, edgecolor="white")
    for b, v in zip(bars, strength):
        axes[1].text(b.get_x() + b.get_width() / 2, v + 0.05, f"{v:.2f}", ha="center", fontsize=10)
    axes[1].set_ylabel("通讯总强度（通讯概率之和）")
    axes[1].set_title("各生态位：通讯总强度", fontsize=11.5)
    axes[1].grid(axis="y", alpha=0.25)

    fig.suptitle("模拟数据演示：不同生态位的通讯「数量」与「强度」比较（CellChat compareInteractions 风格）",
                 fontsize=13.5, fontweight="bold", y=1.0)
    fig.tight_layout(rect=[0, 0, 1, 0.92])
    return savefig(fig, "fig13_comm_barplot.png")

def fig14_diff_comm(data):
    # 气泡图：侵袭前沿 vs 肿瘤核心
    M_f, det_f = comm_prob_matrix(1)
    M_c, det_c = comm_prob_matrix(0)
    fig, axes = plt.subplots(1, 2, figsize=(13.5, 6.2))

    ax = axes[0]
    pairs = LR_PAIRS
    pf = np.array([p for *_, p in [next(d for d in det_f if d[0] == s and d[1] == lig) for s, lig, *_ in pairs]])
    # 更稳：按顺序取
    pf = np.array([comm_prob_matrix(1)[1][i][5] for i in range(len(pairs))])
    pc = np.array([comm_prob_matrix(0)[1][i][5] for i in range(len(pairs))])
    log2fc = np.log2((pf + 1e-3) / (pc + 1e-3))
    ys = np.arange(len(pairs))[::-1]
    sizes = 60 + 900 * pf
    for y, (p, l2, s_, lig) in enumerate(zip(pairs, log2fc, sizes, [p[1] for p in pairs])):
        col = "#C0392B" if l2 > 0 else "#2980B9"
        ax.scatter(0, y, s=s_, color=col, alpha=0.75, edgecolors="white", lw=1)
        ax.scatter(1, y, s=60 + 900 * pc[len(pairs) - 1 - y], color="none",
                   edgecolors="#95A5A6", lw=1.2)
        ax.text(-0.55, y, lig, ha="right", va="center", fontsize=9.5)
    ax.set_yticks([])
    ax.set_xticks([0, 1]); ax.set_xticklabels(["侵袭前沿", "肿瘤核心"], fontsize=10.5)
    ax.set_xlim(-1.2, 1.6)
    ax.set_title("配体-受体对通讯概率（实心=前沿，空心=核心）", fontsize=11)
    ax.grid(axis="y", alpha=0.2)
    ax.text(1.42, len(pairs) - 0.5, "红色：前沿上调\n蓝色：前沿下调", fontsize=9, va="top", color="#34495E")

    ax = axes[1]
    diff = pf - pc
    vmax = np.abs(diff).max()
    im = ax.imshow(diff[:, None], cmap="RdBu_r", vmin=-vmax, vmax=vmax, aspect="auto")
    ax.set_yticks(range(len(pairs)))
    ax.set_yticklabels([f"{p[1]} → {p[3]}" for p in pairs][::-1], fontsize=9)
    ax.set_xticks([])
    for i, d in enumerate(diff[::-1]):
        ax.text(0, i, f"{d:+.2f}", ha="center", va="center", fontsize=9.5,
                color="white" if abs(d) > vmax * 0.55 else "#2C3E50")
    ax.set_title("通讯概率差（前沿 减 核心）", fontsize=11)
    cb = fig.colorbar(im, ax=ax, shrink=0.8); cb.set_label("Δ 通讯概率", fontsize=9)

    fig.suptitle("模拟数据演示：生态位间差异通讯（侵袭前沿 vs 肿瘤核心）——'重编程'候选信号", fontsize=13.5, fontweight="bold", y=1.0)
    fig.tight_layout(rect=[0, 0, 1, 0.93])
    return savefig(fig, "fig14_diff_comm.png")

def fig15_pathway_heatmap(data):
    paths = ["TGFb", "PD-L1", "MIF", "GALECTIN", "VEGF", "CXCL", "ECM", "IL10", "CCL"]
    act = np.zeros((4, len(paths)))
    for k in range(4):
        for j, pw in enumerate(paths):
            act[k, j] = NICHE_PATHWAY_ACT[NICHES[k]][pw]
    # 聚类排序（按前沿 vs 核心差异）
    order = np.argsort(act[1] - act[0])[::-1]
    act = act[:, order]; paths_sorted = [paths[o] for o in order]

    fig, ax = plt.subplots(figsize=(9.6, 5.4))
    im = ax.imshow(act, cmap="YlOrRd", vmin=0, vmax=1)
    ax.set_xticks(range(len(paths_sorted))); ax.set_xticklabels(paths_sorted, rotation=35, ha="right", fontsize=10)
    ax.set_yticks(range(4)); ax.set_yticklabels(NICHES, fontsize=10.5)
    for i in range(4):
        for j in range(len(paths_sorted)):
            ax.text(j, i, f"{act[i, j]:.2f}", ha="center", va="center", fontsize=9,
                    color="white" if act[i, j] > 0.55 else "#2C3E50")
    ax.set_title("模拟数据演示：信号通路活性热图（生态位 × 通路）", fontsize=12.5, fontweight="bold")
    cb = fig.colorbar(im, ax=ax, shrink=0.85); cb.set_label("相对活性（0–1）", fontsize=9.5)
    # 标注"重编程"热点
    ax.annotate("前沿高活性：免疫抑制/侵袭相关通路", xy=(4.5, 1.05), fontsize=10,
                color="#C0392B", ha="center", fontweight="bold")
    return savefig(fig, "fig15_pathway_heatmap.png")

def _ribbon(ax, x0, y0b, y0t, x1, y1b, y1t, color, alpha=0.45):
    """用三次贝塞尔曲线画一条从 (x0, y0b-y0t) 到 (x1, y1b-y1t) 的带状流。"""
    mid = (x0 + x1) / 2
    verts = [(x0, y0b), (mid, y0b), (mid, y1b), (x1, y1b),
             (x1, y1t), (mid, y1t), (mid, y0t), (x0, y0t)]
    codes = [Path.MOVETO, Path.CURVE4, Path.CURVE4, Path.CURVE4,
             Path.LINETO, Path.CURVE4, Path.CURVE4, Path.CURVE4]
    ax.add_patch(mpatches.PathPatch(Path(verts, codes), facecolor=color,
                                    edgecolor="none", alpha=alpha, lw=0))

def fig16_reprogramming(data):
    # 左：生态位；中：通路（按重编程热点）；右：靶细胞类型
    fig, ax = plt.subplots(figsize=(13.5, 7.2))
    ax.set_xlim(0, 13); ax.set_ylim(0, 10); ax.axis("off")
    ax.text(6.5, 9.5, "模拟数据演示：生态位 → 信号通路 → 靶细胞的'通讯重编程'流量图", ha="center",
            fontsize=13.5, fontweight="bold", color="#2C3E50")

    src = NICHES
    mids = ["TGFb", "PD-L1", "MIF", "GALECTIN", "CCL"]
    tgts = ["CAF", "CD8 T", "巨噬细胞", "内皮细胞"]

    total = {k: sum(NICHE_PATHWAY_ACT[k].values()) for k in src}
    # 源 → 通路 流量
    flow_sm = {}
    for k in src:
        for pw in mids:
            flow_sm[(k, pw)] = NICHE_PATHWAY_ACT[k][pw]
    # 通路 → 靶 流量（按 L-R 对的接收者分配）
    flow_mt = {}
    pw_act = {pw: sum(NICHE_PATHWAY_ACT[k][pw] for k in NICHES) for pw in mids}
    for pw in mids:
        tot = sum(p[5] for p in LR_PAIRS if p[4] == pw)
        if tot <= 0:
            continue
        for (s, lig, r, rec, pw2, p) in LR_PAIRS:
            if pw2 == pw:
                flow_mt[(pw, r)] = flow_mt.get((pw, r), 0) + p / tot * pw_act[pw] * 0.9
    # 归一化
    fmax = max(max(flow_sm.values()), max(flow_mt.values()))
    scale = 2.6 / fmax

    x0, x1, x2 = 1.2, 6.2, 11.0
    # 节点 y 布局
    def col_layout(names, x, width=1.15):
        yc = {}
        n = len(names)
        start = 8.0 - (n - 1) * 0.8
        for i, nm in enumerate(names):
            yc[nm] = start + i * 1.6
            ax.add_patch(FancyBboxPatch((x - width / 2, yc[nm] - 0.42), width, 0.84,
                                        boxstyle="round,pad=0.06", linewidth=1.4,
                                        edgecolor="#34495E", facecolor="#EBF5FB"))
            ax.text(x, yc[nm], nm, ha="center", va="center", fontsize=10.5, fontweight="bold", color="#2C3E50")
        return yc

    yc0 = col_layout(src, x0)
    yc1 = col_layout(mids, x1)
    yc2 = col_layout(tgts, x2)

    # 源 → 通路
    for (k, pw), w in flow_sm.items():
        w_ = w * scale * 0.45
        _ribbon(ax, x0 + 0.62, yc0[k] - w_ / 2, yc0[k] + w_ / 2,
                x1 - 0.62, yc1[pw] - w_ / 2, yc1[pw] + w_ / 2,
                NICHE_COLORS[k], alpha=0.32)
    # 通路 → 靶
    for (pw, r), w in flow_mt.items():
        w_ = w * scale * 0.45
        _ribbon(ax, x1 + 0.62, yc1[pw] - w_ / 2, yc1[pw] + w_ / 2,
                x2 - 0.62, yc2[r] - w_ / 2, yc2[r] + w_ / 2,
                "#8E44AD", alpha=0.28)

    ax.text(x0, 8.9, "来源生态位", ha="center", fontsize=10.5, color="#7F8C8D")
    ax.text(x1, 8.9, "重编程信号通路", ha="center", fontsize=10.5, color="#7F8C8D")
    ax.text(x2, 8.9, "靶细胞类型", ha="center", fontsize=10.5, color="#7F8C8D")
    ax.text(6.5, 0.4, "带宽 ∝ 通讯活性：侵袭前沿与免疫浸润区贡献了大部分免疫抑制/趋化信号", ha="center",
            fontsize=10, color="#34495E", style="italic")
    return savefig(fig, "fig16_reprogramming.png")

def fig17_paper_figure(data):
    fig = plt.figure(figsize=(15, 11))
    gs = fig.add_gridspec(2, 2, hspace=0.18, wspace=0.10)

    # A: 生态位
    ax = fig.add_subplot(gs[0, 0])
    cols = [NICHE_COLORS[NICHES[k]] for k in data["niche"]]
    ax.scatter(data["X"], data["Y"], c=cols, s=6, edgecolors="none")
    ax.set_aspect("equal"); ax.set_xticks([]); ax.set_yticks([])
    ax.set_title("A  空间生态位识别", fontsize=12.5, fontweight="bold", loc="left")

    # B: 邻域富集
    z = _nhood_enrichment(data)
    ax = fig.add_subplot(gs[0, 1])
    vmax = np.abs(z).max() * 0.8
    im = ax.imshow(z, cmap="RdBu_r", vmin=-vmax, vmax=vmax)
    ax.set_xticks(range(6)); ax.set_xticklabels(CTS, rotation=40, ha="right", fontsize=8)
    ax.set_yticks(range(6)); ax.set_yticklabels(CTS, fontsize=8)
    ax.set_title("B  细胞共定位富集", fontsize=12.5, fontweight="bold", loc="left")
    fig.colorbar(im, ax=ax, shrink=0.9)

    # C: 差异通讯气泡
    ax = fig.add_subplot(gs[1, 0])
    pf = np.array([comm_prob_matrix(1)[1][i][5] for i in range(len(LR_PAIRS))])
    pc = np.array([comm_prob_matrix(0)[1][i][5] for i in range(len(LR_PAIRS))])
    l2 = np.log2((pf + 1e-3) / (pc + 1e-3))
    ys = np.arange(len(LR_PAIRS))[::-1]
    for y, (p, l2v, pv) in enumerate(zip(LR_PAIRS, l2, pf)):
        ax.scatter(0, y, s=80 + 700 * pv, color="#C0392B" if l2v > 0 else "#2980B9",
                   alpha=0.8, edgecolors="white")
        ax.scatter(1, y, s=80 + 700 * pc[len(pc) - 1 - y], color="none", edgecolors="#95A5A6")
    ax.set_yticks(range(len(LR_PAIRS)))
    ax.set_yticklabels([p[1] for p in LR_PAIRS][::-1], fontsize=8.5)
    ax.set_xticks([0, 1]); ax.set_xticklabels(["侵袭前沿", "肿瘤核心"], fontsize=9)
    ax.set_xlim(-0.8, 1.8)
    ax.set_title("C  生态位间差异通讯（配体-受体）", fontsize=12.5, fontweight="bold", loc="left")

    # D: 通路热图
    ax = fig.add_subplot(gs[1, 1])
    paths = ["TGFb", "PD-L1", "MIF", "GALECTIN", "VEGF", "CXCL", "ECM", "IL10", "CCL"]
    act = np.array([[NICHE_PATHWAY_ACT[NICHES[k]][pw] for pw in paths] for k in range(4)])
    im = ax.imshow(act, cmap="YlOrRd", vmin=0, vmax=1)
    ax.set_xticks(range(len(paths))); ax.set_xticklabels(paths, rotation=35, ha="right", fontsize=8)
    ax.set_yticks(range(4)); ax.set_yticklabels(NICHES, fontsize=9)
    ax.set_title("D  通路活性（生态位 × 通路）", fontsize=12.5, fontweight="bold", loc="left")
    fig.colorbar(im, ax=ax, shrink=0.9)

    fig.suptitle("模拟数据演示：论文结果图版式示例（Figure 2 风格：生态位 → 共定位 → 差异通讯 → 通路活性）",
                 fontsize=14, fontweight="bold", y=0.99)
    return savefig(fig, "fig17_paper_figure.png")

# ============================================================ main
def main():
    data = sim_visium()
    # 保存模拟数据集（复现用）
    os.makedirs(DATA_DIR, exist_ok=True)
    np.savez(os.path.join(DATA_DIR, "sim_visium.npz"),
             X=data["X"], Y=data["Y"], niche=data["niche"], prop=data["prop"],
             counts=data["counts"], nCounts=data["nCounts"], nGenes=data["nGenes"],
             pct_mito=data["pct_mito"], cluster=data["cluster"], lowq=data["lowq"])
    n = len(data["X"])
    print(f"[sim] spots={n} 生态位分布=", dict(zip(NICHES, np.bincount(data["niche"], minlength=4))))

    fig07_qc(data)
    fig08_clustering(data)
    fig09_celltype(data)
    fig10_niche(data)
    fig11_nhood(data)
    fig12_comm_network(data)
    fig13_comm_barplot(data)
    fig14_diff_comm(data)
    fig15_pathway_heatmap(data)
    fig16_reprogramming(data)
    fig17_paper_figure(data)
    print("Part B done.")

if __name__ == "__main__":
    main()
