# =========================================================
# Academic-style Ablation Figure
# =========================================================

import matplotlib.pyplot as plt
import numpy as np

configs = [
    "Biomni\n(Baseline)",
    "+MCP\nInfra.",
    "+MCP Infra.\n +Graph-Sca\n Planning",
]

# 仅更新此处的数字，完全对齐表格中的最新数据
metrics = {
    "BioAgentBench": [39.67, 42.68, 46.84],
    "LAB-Bench DbQA": [57.46, 66.35, 67.29],
    "LAB-Bench SeqQA": [84.76, 89.84, 90.48],
}

# =========================================================
# Muted academic palette
# =========================================================
fill_colors = [
    "#E5E7EB",   # soft gray
    "#D6E4F0",   # muted blue
    "#E8D6DC",   # muted wine
]

edge_colors = [
    "#6B7280",
    "#4B6A88",
    "#8A3D5D",
]

# =========================================================
# Style
# =========================================================
plt.rcParams.update(
    {
        "font.family": "DejaVu Sans",
        "font.size": 10,
        "axes.titlesize": 12,
        "xtick.labelsize": 9,
        "ytick.labelsize": 9.5,
        "pdf.fonttype": 42,
        "ps.fonttype": 42,
    }
)

fig, axes = plt.subplots(
    1,
    3,
    figsize=(8.4, 2.85),
    sharey=True,
)

y = np.arange(len(configs))

for idx, (ax, (title, values)) in enumerate(zip(axes, metrics.items())):

    values = np.array(values)
    baseline = values[0]

    span = values.max() - values.min()

    x_min = values.min() - span * 0.28 - 0.6
    x_max = values.max() + span * 0.52 + 1.2

    ax.set_xlim(x_min, x_max)

    # subtle baseline reference
    ax.axvline(
        baseline,
        color="#C7CCD4",
        linestyle="--",
        linewidth=0.9,
        zorder=1,
    )

    for i, val in enumerate(values):

        # main bar
        ax.barh(
            y[i],
            val - x_min,
            left=x_min,
            height=0.42,
            color=fill_colors[i],
            edgecolor=edge_colors[i],
            linewidth=1.15,
            zorder=3,
        )

        # subtle inner highlight
        ax.barh(
            y[i],
            (val - x_min) * 0.52,
            left=x_min,
            height=0.42,
            color="white",
            alpha=0.12,
            linewidth=0,
            zorder=4,
        )

        # value inside bar
        ax.text(
            val - span * 0.06,
            y[i],
            f"{val:.1f}",
            ha="right",
            va="center",
            fontsize=10,
            color=edge_colors[i],
            fontweight="bold" if i == 2 else "normal",
        )

        # gain
        if i > 0:
            gain = val - baseline

            ax.text(
                val + span * 0.14 + 0.2,
                y[i],
                f"+{gain:.1f}",
                ha="left",
                va="center",
                fontsize=8.8,
                color="#16924A",
            )

    # panel title
    ax.set_title(title, pad=7, fontweight="bold")

    ax.set_xticks([])

    # minimal style
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.spines["bottom"].set_visible(False)

    ax.spines["left"].set_color("#D1D5DB")
    ax.spines["left"].set_linewidth(0.8)

    if idx == 0:
        ax.set_yticks(y)
        ax.set_yticklabels(configs)
    else:
        ax.tick_params(axis="y", left=False, labelleft=False)

# top-to-bottom
axes[0].invert_yaxis()

# =========================================================
# Figure title
# =========================================================
fig.suptitle(
    "Ablation of MCP Infrastructure and Graph Planning",
    fontsize=13.5,
    fontweight="bold",
    y=1.02,
)

# =========================================================
# Layout
# =========================================================
fig.subplots_adjust(
    left=0.16,
    right=0.985,
    top=0.72,
    bottom=0.15,
    wspace=0.08,
)

# =========================================================
# Save
# =========================================================
fig.savefig("ablation_academic.pdf", bbox_inches="tight")
fig.savefig("ablation_academic.png", dpi=500, bbox_inches="tight")
fig.savefig("ablation_academic.svg", bbox_inches="tight")