# -*- coding: utf-8 -*-
"""
Reproduction of Fig. 2b from:
    Zika virus vaccine with E protein fusion loop mutations protects via CD8+ T cells
    Nature Microbiology 2026, doi:10.1038/s41564-026-02465-6
Panel: serum anti-ZIKV E IgG (endpoint ELISA, log10) and NT50 (50% neutralizing titer, log10),
       CON / WT / FLM groups.

数据说明（模拟数据！）：
    本论文未随刊发布该图的 Source Data（PMC13652353 无附件源数据）。
    以下数据为作者手工按原图散点分布手工估读并生成、固定种子(seed=20261019)的**模拟数据**，
    仅用于教学复现作图代码，数值均值/离散度/样本量与原图大致相符（WT n=16, FLM n=18, CON n=6），
    不代表论文原始数据。图中的显著性 p 值引用自论文原文。

Tested: python3 (matplotlib) — sandbox run OK.
"""
import os
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib import font_manager

HERE = os.path.dirname(os.path.abspath(__file__))
FONT = os.path.join(HERE, "..", "day1-figure", "NotoSansSC-Regular.ttf")
font_manager.fontManager.addfont(os.path.normpath(FONT))
plt.rcParams["font.family"] = "Noto Sans SC"
plt.rcParams["axes.unicode_minus"] = False

RNG = np.random.default_rng(20261019)

# ---------------- 模拟数据（seed 固定） ----------------
spec = {
    # (n, mean, sd, clamp_low, clamp_high) —— 模拟参数按原图散点范围手工设定
    "igg":   {"CON": (6, 1.46, 0.03, 1.40, 1.52),
              "WT": (16, 3.58, 0.40, 2.70, 4.60),
              "FLM": (18, 3.52, 0.42, 2.40, 4.65)},
    "nt50":  {"CON": (6, 0.97, 0.02, 0.92, 1.02),
              "WT": (16, 2.36, 0.30, 1.70, 3.00),
              "FLM": (18, 2.78, 0.24, 2.20, 3.30)},
}
rows = []
for panel, groups in spec.items():
    for g, (n, mean, sd, lo, hi) in groups.items():
        vals = np.clip(RNG.normal(mean, sd, n), lo, hi)
        for i, v in enumerate(vals):
            rows.append({"panel": panel, "group": g, "value_log10": round(float(v), 3),
                         "point_type": "individual"})
df = pd.DataFrame(rows)

# 模拟原图的红色点（pooled sera）：取各组最接近组中位数的点标红
for (panel, g, n_red) in [("igg", "WT", 2), ("igg", "FLM", 2), ("igg", "CON", 1),
                          ("nt50", "WT", 2), ("nt50", "FLM", 1), ("nt50", "CON", 2)]:
    sel_g = df[(df.panel == panel) & (df.group == g)]
    med = sel_g.value_log10.median()
    sel = (sel_g.value_log10 - med).abs().nsmallest(n_red).index
    df.loc[sel, "point_type"] = "pooled"
df.to_csv(os.path.join(HERE, "zika_vaccine_fig2b.csv"), index=False)

# ---------------- 作图 ----------------
groups = ["CON", "WT", "FLM"]
bar_fc = {"CON": "#FFFFFF", "WT": "#000000", "FLM": "#000000"}
p_texts = {
    "igg":  [("CON", "WT", 5.10, 4.82, "<0.0001"), ("CON", "FLM", 5.60, 5.32, "<0.0001")],
    "nt50": [("CON", "WT", 3.95, 3.68, "0.0010"),  ("CON", "FLM", 4.45, 4.18, "<0.0001")],
}
lod = {"igg": 1.45, "nt50": 1.00}
titles = {"igg": "Anti-ZIKV E IgG (log$_{10}$)", "nt50": "NT50 (log$_{10}$)"}

fig, axes = plt.subplots(1, 2, figsize=(8.4, 4.4), dpi=300)
axes[0].text(-0.16, 1.06, "b", transform=axes[0].transAxes,
             fontsize=20, fontweight="bold")
for ax, panel in zip(axes, ["igg", "nt50"]):
    sub = df[df.panel == panel]
    x = np.arange(len(groups))
    means = [sub[sub.group == g].value_log10.mean() for g in groups]
    sems = [sub[sub.group == g].value_log10.std(ddof=1) / np.sqrt(sub[sub.group == g].shape[0]) for g in groups]
    ax.bar(x, means, width=0.72, color=[bar_fc[g] for g in groups],
           edgecolor="#000000", linewidth=1.2, zorder=2)
    ax.errorbar(x, means, yerr=sems, fmt="none", ecolor="#000000",
                elinewidth=1.2, capsize=5, capthick=1.2, zorder=3)
    for xi, g in enumerate(groups):
        pts = sub[sub.group == g]
        for _, r in pts.iterrows():
            jitter = RNG.uniform(-0.22, 0.22)
            fc = "#D32F2F" if r.point_type == "pooled" else "#FFFFFF"
            ax.scatter(xi + jitter, r.value_log10, s=42, c=fc,
                       edgecolors="#000000", linewidths=0.9, zorder=4)
    ax.axhline(lod[panel], color="#000000", linestyle=(0, (2, 2)), linewidth=1.0, alpha=0.7)
    for (g1, g2, y_top, y_txt, p) in p_texts[panel]:
        x1, x2 = groups.index(g1), groups.index(g2)
        ax.plot([x1, x1, x2, x2], [y_txt + 0.18, y_top, y_top, y_txt + 0.18],
                color="#000000", linewidth=1.1, clip_on=False)
        ax.text((x1 + x2) / 2, y_top + 0.05, p, ha="center", va="bottom", fontsize=8.5)
    ax.set_xticks(x)
    ax.set_xticklabels(groups, rotation=45, ha="right", fontsize=10)
    ax.set_ylabel(titles[panel], fontsize=10.5)
    ax.set_ylim(0, 5.9 if panel == "igg" else 5.0)
    ax.set_yticks(range(0, 6))
    ax.tick_params(axis="y", which="both", direction="out", length=4)
    ax.minorticks_on()
    for spine in ("top", "right"):
        ax.spines[spine].set_visible(False)
fig.tight_layout(rect=[0.02, 0.03, 0.99, 0.93])
out = os.path.join(HERE, "reproduce_py.png")
fig.savefig(out, dpi=300)
print("saved:", out)