#!/usr/bin/env python3
from __future__ import annotations

import re
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd


# =========================
# Config
# =========================
XLSX_PATH = Path("Context_Consuming.xlsx")
OUT_PREFIX = Path("bioagentbench_context_radar_final")

METRIC = "prompt_tokens"

# 正文图建议只保留这 4 个；如果要完整 6 条线，见下方注释。
# METHOD_ORDER = [
#     "No MCP",
#     "Biomni-100",
#     "Biomni-2k",
#     "BioManus",
# ]

# 附录完整版本可改成：
METHOD_ORDER = [
    "No MCP",
    "Biomni-100",
    "Biomni-500",
    "Biomni-1k",
    "Biomni-2k",
    "BioManus",
]

TASK_ORDER = [
    "alzheimer-mouse",
    "cystic-fibrosis",
    "deseq",
    "transcript-quant",
    "single-cell",
    "metagenomics",
    "viral-metagenomics",
    "comparative-genomics",
    "evolution",
    "giab",
]

METHOD_STYLE = {
    "No MCP": {
        "color": "#6B7280",
        "linewidth": 1.8,
        "linestyle": "-",
        "fill_alpha": 0.00,
    },
    "Biomni-100": {
        "color": "#2563EB",
        "linewidth": 1.8,
        "linestyle": "-",
        "fill_alpha": 0.00,
    },
    "Biomni-500": {
        "color": "#059669",
        "linewidth": 1.6,
        "linestyle": "-",
        "fill_alpha": 0.00,
    },
    "Biomni-1k": {
        "color": "#D97706",
        "linewidth": 1.6,
        "linestyle": "-",
        "fill_alpha": 0.00,
    },
    "Biomni-2k": {
        "color": "#7C3AED",
        "linewidth": 2.0,
        "linestyle": "-",
        "fill_alpha": 0.00,
    },
    "BioManus": {
        "color": "#DC2626",
        "linewidth": 2.6,
        "linestyle": "-",
        "fill_alpha": 0.04,
    },
}


# =========================
# Utilities
# =========================
def clean_text(x) -> str:
    if pd.isna(x):
        return ""
    s = str(x)
    s = s.replace("\ufeff", "").replace("\u200b", "").replace("\u200c", "")
    s = re.sub(r"\s+", " ", s)
    return s.strip()


def method_name(x) -> str | None:
    s = clean_text(x)
    if not s:
        return None

    lower = s.lower()
    if "biomanus" in lower:
        return "BioManus"

    try:
        scale = int(float(s))
    except ValueError:
        return None

    if scale == 0:
        return "No MCP"
    if scale == 100:
        return "Biomni-100"
    if scale == 500:
        return "Biomni-500"
    if scale == 1000:
        return "Biomni-1k"
    if scale == 2000:
        return "Biomni-2k"

    # Ignore 250 and other scales.
    return None


def wrap_task_label(task: str) -> str:
    labels = {
        "alzheimer-mouse": "alzheimer\nmouse",
        "cystic-fibrosis": "cystic\nfibrosis",
        "deseq": "deseq",
        "transcript-quant": "transcript\nquant",
        "single-cell": "single-cell",
        "metagenomics": "metagenomics",
        "viral-metagenomics": "viral\nmetagenomics",
        "comparative-genomics": "comparative\ngenomics",
        "evolution": "evolution",
        "giab": "giab",
    }
    return labels.get(task, task.replace("-", "\n"))


def geometric_mean(values: np.ndarray) -> float:
    values = np.asarray(values, dtype=float)
    values = values[values > 0]
    return float(np.exp(np.mean(np.log(values)))) if len(values) else np.nan


def format_token(x: float) -> str:
    if x >= 1_000_000:
        return f"{x / 1_000_000:.1f}M"
    if x >= 1_000:
        return f"{x / 1_000:.0f}K"
    return f"{x:.0f}"


# =========================
# Load data
# =========================
raw = pd.read_excel(XLSX_PATH, sheet_name=0, header=None)

header_candidates = raw.index[
    raw.iloc[:, 1].astype(str).map(clean_text).eq("Tasks")
].tolist()

if not header_candidates:
    raise ValueError("Cannot find header row with 'Tasks' in column B.")

header_idx = header_candidates[0]

df = raw.iloc[header_idx + 1 :, :6].copy()
df.columns = [
    "group",
    "task",
    "results_match",
    "prompt_tokens",
    "completion_tokens",
    "total_tokens",
]

df["group"] = df["group"].ffill()
df["method"] = df["group"].map(method_name)
df["task"] = df["task"].map(clean_text)

df = df[df["task"].ne("")]
df = df[df["method"].notna()]

for col in ["prompt_tokens", "completion_tokens", "total_tokens"]:
    df[col] = pd.to_numeric(df[col], errors="coerce")

df = df.dropna(subset=[METRIC])
df = df[df["method"].isin(METHOD_ORDER)]
df["method"] = pd.Categorical(df["method"], categories=METHOD_ORDER, ordered=True)

pivot = df.pivot_table(
    index="method",
    columns="task",
    values=METRIC,
    aggfunc="mean",
)

available_methods = [m for m in METHOD_ORDER if m in pivot.index]
task_order = [t for t in TASK_ORDER if t in pivot.columns]

common_tasks = [
    t for t in task_order
    if pivot.loc[available_methods, t].notna().all()
]

pivot = pivot.loc[available_methods, common_tasks]

if pivot.empty:
    raise ValueError("No valid method-task matrix found for radar plotting.")


# =========================
# Plot
# =========================
plt.rcParams.update(
    {
        "font.family": "DejaVu Sans",
        "font.size": 10,
        "axes.titlesize": 14,
        "legend.fontsize": 9,
        "pdf.fonttype": 42,
        "ps.fonttype": 42,
    }
)

tasks = list(pivot.columns)
task_labels = [wrap_task_label(t) for t in tasks]
n_tasks = len(tasks)

angles = np.linspace(0, 2 * np.pi, n_tasks, endpoint=False)
angles_closed = np.concatenate([angles, [angles[0]]])

fig, ax = plt.subplots(
    figsize=(7.6, 6.6),
    subplot_kw={"projection": "polar"},
)

ax.set_theta_offset(np.pi / 2)
ax.set_theta_direction(-1)

values_log = np.log10(pivot.values.astype(float))

# Fixed range keeps the figure visually stable across reruns.
r_min = 5.0
r_max = max(7.2, float(np.nanmax(values_log)) + 0.15)
ax.set_ylim(r_min, r_max)

# Axis labels.
ax.set_xticks(angles)
ax.set_xticklabels(task_labels, fontsize=9.5)

# Radial ticks.
radial_ticks = [5, 6, 7]
radial_ticks = [t for t in radial_ticks if r_min <= t <= r_max]
ax.set_yticks(radial_ticks)
ax.set_yticklabels([rf"$10^{t}$" for t in radial_ticks], fontsize=9, color="#4B5563")
ax.set_rlabel_position(90)

# Grid style.
ax.grid(True, color="#CBD5E1", linewidth=0.75, alpha=0.75)
ax.spines["polar"].set_color("#94A3B8")
ax.spines["polar"].set_linewidth(0.8)

# Plot methods.
for method in available_methods:
    values = np.log10(pivot.loc[method, tasks].values.astype(float))
    values_closed = np.concatenate([values, [values[0]]])
    style = METHOD_STYLE[method]

    ax.plot(
        angles_closed,
        values_closed,
        color=style["color"],
        linewidth=style["linewidth"],
        linestyle=style["linestyle"],
        label=method,
        zorder=4 if method == "BioManus" else 3,
    )

    if style["fill_alpha"] > 0:
        ax.fill(
            angles_closed,
            values_closed,
            color=style["color"],
            alpha=style["fill_alpha"],
            zorder=2,
        )

# Title.
ax.set_title("BioAgentBench Context Consumption", pad=22, fontweight="bold")

# Legend.
ax.legend(
    loc="upper left",
    bbox_to_anchor=(1.05, 1.03),
    frameon=False,
    handlelength=2.4,
    borderaxespad=0.0,
)

fig.subplots_adjust(left=0.05, right=0.78, top=0.90, bottom=0.06)

# Save.
fig.savefig(f"{OUT_PREFIX}.pdf", bbox_inches="tight")
fig.savefig(f"{OUT_PREFIX}.png", dpi=400, bbox_inches="tight")
fig.savefig(f"{OUT_PREFIX}.svg", bbox_inches="tight")

print("Saved:")
print(f"  {OUT_PREFIX}.pdf")
print(f"  {OUT_PREFIX}.png")
print(f"  {OUT_PREFIX}.svg")

print("\nGeometric mean prompt tokens:")
for method in available_methods:
    gm = geometric_mean(pivot.loc[method, tasks].values.astype(float))
    print(f"{method:12s}: {format_token(gm)}")