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

from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.ticker import FuncFormatter


# ============================================================
# Load Excel
# ============================================================

SCRIPT_DIR = Path(__file__).resolve().parent
xlsx_path = SCRIPT_DIR / "Context_Consuming.xlsx"

# BioAgentBench table
bio_df = pd.read_excel(
    xlsx_path,
    sheet_name=0,
    usecols="A:E",
    header=1,
)

bio_df.columns = [
    "scale",
    "task",
    "prompt",
    "completion",
    "total",
]

bio_df["scale"] = bio_df["scale"].ffill()
bio_df["task"] = (
    bio_df["task"]
    .astype(str)
    .str.replace("\ufeff", "", regex=False)
    .str.strip()
)

scale_map = {
    0: "No MCP",
    100: "Biomni-100",
    500: "Biomni-500",
    1000: "Biomni-1k",
    2000: "Biomni-2k",
}

bio_df["method"] = bio_df["scale"].map(scale_map)

# Last 10 rows are BioManus in your table
bio_last = bio_df.tail(10).copy()
bio_last["method"] = "BioManus"
bio_df = pd.concat([bio_df.iloc[:-10], bio_last], ignore_index=True)

bio_df = bio_df[bio_df["method"].notna()].copy()


# Closing Scenarios table
closing_df = pd.read_excel(
    xlsx_path,
    sheet_name=0,
    usecols="I:L",
    header=1,
)

closing_df.columns = [
    "model",
    "prompt",
    "completion",
    "total",
]

closing_df["model"] = closing_df["model"].astype(str).str.strip()

closing_name_map = {
    "Biomni": "Biomni",
    "Biomni with 100 tool scale": "Biomni-100",
    "Biomni with 500 tool scale": "Biomni-500",
    "Biomni with 1000 tool scale": "Biomni-1k",
    "Biomni with 2000 tool scale": "Biomni-2k",
    "Biomanus": "BioManus",
}

closing_df["method"] = closing_df["model"].map(closing_name_map)
closing_df = closing_df[closing_df["method"].notna()].copy()


# ============================================================
# Orders and labels
# ============================================================

method_order = [
    "No MCP",
    "Biomni-100",
    "Biomni-500",
    "Biomni-1k",
    "Biomni-2k",
    "BioManus",
]

closing_order = [
    "Biomni",
    "Biomni-100",
    "Biomni-500",
    "Biomni-1k",
    "Biomni-2k",
    "BioManus",
]

task_order = (
    bio_df[bio_df["method"] == "No MCP"]["task"]
    .drop_duplicates()
    .tolist()
)

# Optional: reorder tasks by average context cost from low to high
task_score = (
    bio_df.groupby("task")["prompt"]
    .mean()
    .sort_values()
)

task_order = [t for t in task_score.index.tolist() if t in task_order]


def wrap_task_label(task: str) -> str:
    mapping = {
        "alzheimer-mouse": "Alzheimer\nmouse",
        "comparative-genomics": "Comparative\ngenomics",
        "cystic-fibrosis": "Cystic\nfibrosis",
        "deseq": "DESeq",
        "evolution": "Evolution",
        "giab": "GIAB",
        "metagenomics": "Metagenomics",
        "single-cell": "Single-cell",
        "transcript-quant": "Transcript\nquant",
        "viral-metagenomics": "Viral\nmetagenomics",
    }
    return mapping.get(task, task.replace("-", "\n"))


# ============================================================
# Style
# ============================================================

colors = {
    "No MCP": "#8A8F98",
    "Biomni": "#8A8F98",
    "Biomni-100": "#C9853B",
    "Biomni-500": "#B8742A",
    "Biomni-1k": "#9C5C1A",
    "Biomni-2k": "#6F3B0D",
    "BioManus": "#C62828",
}

fills = {
    "No MCP": "#ECEFF3",
    "Biomni": "#ECEFF3",
    "Biomni-100": "#F1D9BA",
    "Biomni-500": "#E8C390",
    "Biomni-1k": "#D8A96E",
    "Biomni-2k": "#C69054",
    "BioManus": "#F1C9C9",
}

plt.rcParams.update(
    {
        "font.family": "DejaVu Sans",
        "font.size": 10.3,
        "axes.titlesize": 12.3,
        "axes.labelsize": 10.3,
        "xtick.labelsize": 8.8,
        "ytick.labelsize": 9.0,
        "legend.fontsize": 8.6,
        "pdf.fonttype": 42,
        "ps.fonttype": 42,
    }
)


# ============================================================
# Helpers
# ============================================================

def token_formatter(value: float, _pos=None):
    if value >= 1_000_000:
        return f"{value / 1_000_000:.1f}M"
    if value >= 1_000:
        return f"{value / 1_000:.0f}K"
    return f"{value:.0f}"


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


def percent_reduction(base: float, target: float) -> float:
    return 100.0 * (base - target) / base


# ============================================================
# Prepare Panel A matrix
# ============================================================

pivot = bio_df.pivot_table(
    index="method",
    columns="task",
    values="prompt",
    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 and pivot.loc[available_methods, t].notna().all()
]

pivot = pivot.loc[available_methods, task_order]

log_pivot = np.log10(pivot.astype(float))

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


# ============================================================
# Figure
# ============================================================

fig = plt.figure(figsize=(12.2, 5.65))

gs = fig.add_gridspec(
    1,
    2,
    width_ratios=[1.25, 1.0],
    wspace=0.42,
)

ax_a = fig.add_subplot(gs[0, 0], projection="polar")
ax_b = fig.add_subplot(gs[0, 1])


# ============================================================
# Panel A: Radar chart
# ============================================================

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

r_min = 5.0
r_data_max = max(7.3, float(np.nanmax(log_pivot.values)) + 0.15)
r_max = r_data_max + 0.18
ax_a.set_ylim(r_min, r_max)

ax_a.set_xticks(angles)
ax_a.set_xticklabels([])

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

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

label_radius = r_max + 0.34
for angle, task in zip(angles, task_order):
    display_angle = np.pi / 2 - angle
    x = np.cos(display_angle)
    y_pos = np.sin(display_angle)
    ha = "center"
    if x > 0.22:
        ha = "left"
    elif x < -0.22:
        ha = "right"
    va = "center"
    if y_pos > 0.78:
        va = "bottom"
    elif y_pos < -0.78:
        va = "top"
    ax_a.text(
        angle,
        label_radius,
        wrap_task_label(task),
        ha=ha,
        va=va,
        fontsize=8.2,
        color="#374151",
        clip_on=False,
        bbox={
            "boxstyle": "round,pad=0.12",
            "facecolor": "white",
            "edgecolor": "none",
            "alpha": 0.82,
        },
    )

for method in available_methods:
    values = log_pivot.loc[method, task_order].values.astype(float)
    values_closed = np.concatenate([values, [values[0]]])

    lw = 2.5 if method == "BioManus" else 1.45
    alpha = 1.0 if method == "BioManus" else 0.82

    ax_a.plot(
        angles_closed,
        values_closed,
        color=colors[method],
        linewidth=lw,
        alpha=alpha,
        marker="o",
        markersize=4.8 if method == "BioManus" else 3.2,
        markeredgecolor="white",
        markeredgewidth=0.55,
        label=method,
    )

    if method == "BioManus":
        ax_a.fill(
            angles_closed,
            values_closed,
            color=colors[method],
            alpha=0.06,
        )

        raw_values = pivot.loc[method, task_order].values.astype(float)
        for angle, radius, raw_value in zip(angles, values, raw_values):
            display_angle = np.pi / 2 - angle
            horizontal = np.cos(display_angle)
            vertical = np.sin(display_angle)
            text_radius = min(radius + 0.13, r_max - 0.04)
            ha = "left" if horizontal > 0.18 else "right" if horizontal < -0.18 else "center"
            va = "bottom" if vertical > 0.18 else "top" if vertical < -0.18 else "center"
            ax_a.text(
                angle,
                text_radius,
                compact_token_formatter(raw_value),
                ha=ha,
                va=va,
                fontsize=6.9,
                color=colors[method],
                fontweight="bold",
                clip_on=False,
                bbox={
                    "boxstyle": "round,pad=0.10",
                    "facecolor": "white",
                    "edgecolor": "none",
                    "alpha": 0.72,
                },
            )

ax_a.set_title(
    "",
)

ax_a.legend(
    loc="upper center",
    bbox_to_anchor=(0.5, -0.20),
    ncol=3,
    frameon=False,
    handlelength=2.2,
)


# ============================================================
# Panel B: Closing Scenarios aggregate comparison
# ============================================================

closing_values = np.asarray([
    closing_df[closing_df["method"] == m]["prompt"].iloc[0]
    for m in closing_order
])

y = np.arange(len(closing_order))

for i, method in enumerate(closing_order):
    ax_b.barh(
        y[i],
        closing_values[i],
        height=0.48,
        color=fills[method],
        edgecolor=colors[method],
        linewidth=1.25,
        zorder=3,
    )

    ax_b.scatter(
        closing_values[i],
        y[i],
        s=32,
        color=colors[method],
        edgecolor="white",
        linewidth=0.6,
        zorder=4,
    )

    ax_b.text(
        closing_values[i] + 45_000,
        y[i],
        token_formatter(closing_values[i]),
        va="center",
        ha="left",
        fontsize=8.8,
        color=colors[method],
        fontweight="bold" if method == "BioManus" else "normal",
    )

ax_b.set_yticks(y)
ax_b.set_yticklabels(closing_order)
ax_b.invert_yaxis()

ax_b.xaxis.set_major_formatter(FuncFormatter(token_formatter))
ax_b.set_xlabel("Average prompt tokens per case")
ax_b.set_title(
    "",
)

ax_b.grid(axis="x", linestyle="-", linewidth=0.6, alpha=0.22)
ax_b.spines["top"].set_visible(False)
ax_b.spines["right"].set_visible(False)
ax_b.spines["left"].set_color("#CBD5E1")
ax_b.spines["bottom"].set_color("#CBD5E1")

ax_b.set_xlim(0, max(closing_values) * 1.22)

bio = closing_df[closing_df["method"] == "BioManus"]["prompt"].iloc[0]
biomni = closing_df[closing_df["method"] == "Biomni"]["prompt"].iloc[0]
biomni2k = closing_df[closing_df["method"] == "Biomni-2k"]["prompt"].iloc[0]

ax_b.text(
    0.03,
    1.04,
    f"BioManus reduces context\n"
    f"{percent_reduction(biomni, bio):.1f}% vs. Biomni\n"
    f"{percent_reduction(biomni2k, bio):.1f}% vs. Biomni-2k",
    transform=ax_b.transAxes,
    fontsize=8.8,
    color="#374151",
    va="bottom",
    bbox={
        "boxstyle": "round,pad=0.35",
        "facecolor": "white",
        "edgecolor": "#E5E7EB",
        "alpha": 0.95,
    },
    clip_on=False,
)

fig.text(
    0.028,
    0.948,
    "(a) BioAgentBench task-level context profile",
    ha="left",
    va="top",
    fontsize=12.3,
    fontweight="bold",
)
fig.text(
    0.635,
    0.948,
    "(b) Closing Scenarios aggregate context",
    ha="left",
    va="top",
    fontsize=12.3,
    fontweight="bold",
)


# ============================================================
# Save
# ============================================================
fig.subplots_adjust(
    left=0.075,
    right=0.985,
    top=0.76,
    bottom=0.24,
    wspace=0.42,
)
out = SCRIPT_DIR / "context_efficiency_radar_bar"
fig.savefig(f"{out}.pdf", bbox_inches="tight")
fig.savefig(f"{out}.svg", bbox_inches="tight")
fig.savefig(f"{out}.png", dpi=450, bbox_inches="tight")

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