#!/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_parallel_warm")

METRIC = "prompt_tokens"

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

METHOD_STYLE = {
    "No MCP": {
        "color": "#8A8F98",
        "linewidth": 1.7,
        "alpha": 0.72,
        "zorder": 2,
    },
    "Biomni-100": {
        "color": "#C77C2A",
        "linewidth": 1.8,
        "alpha": 0.78,
        "zorder": 3,
    },
    "Biomni-500": {
        "color": "#B66D18",
        "linewidth": 1.8,
        "alpha": 0.80,
        "zorder": 3,
    },
    "Biomni-1k": {
        "color": "#9A5200",
        "linewidth": 1.9,
        "alpha": 0.86,
        "zorder": 4,
    },
    "Biomni-2k": {
        "color": "#7A3E00",
        "linewidth": 2.1,
        "alpha": 0.92,
        "zorder": 5,
    },
    "BioManus": {
        "color": "#D62728",
        "linewidth": 2.8,
        "alpha": 1.00,
        "zorder": 7,
    },
}


# =========================================================
# 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_label(task: str) -> str:
    mapping = {
        "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 mapping.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 Excel
# =========================================================
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 containing '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]

common_tasks = [
    t for t in pivot.columns
    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.")

# Sort tasks from low to high context.
log_pivot = np.log10(pivot.astype(float))
task_score = log_pivot.mean(axis=0).sort_values(ascending=True)
tasks = task_score.index.tolist()
log_pivot = log_pivot.loc[available_methods, tasks]


# =========================================================
# Plot
# =========================================================
plt.rcParams.update(
    {
        "font.family": "DejaVu Serif",
        "font.size": 10.5,
        "axes.titlesize": 14.5,
        "axes.labelsize": 11,
        "legend.fontsize": 9.2,
        "xtick.labelsize": 9.0,
        "ytick.labelsize": 9.5,
        "pdf.fonttype": 42,
        "ps.fonttype": 42,
    }
)

fig, ax = plt.subplots(figsize=(10.2, 4.8))

x = np.arange(len(tasks))
y_min = 5.0
y_max = max(7.2, float(np.nanmax(log_pivot.values)) + 0.15)

# Warm background band similar to the reference style.
ax.axhspan(y_min, y_max, color="#F6EFE7", alpha=0.55, zorder=0)

# Vertical translucent bars at each task axis.
for xi in x:
    ax.vlines(
        xi,
        y_min,
        y_max,
        color="#D99A5B",
        linewidth=5.0,
        alpha=0.16,
        zorder=1,
    )

# Horizontal subtle grid.
for y in [5, 6, 7]:
    if y_min <= y <= y_max:
        ax.axhline(
            y,
            color="#C9C1B8",
            linewidth=0.8,
            linestyle=":",
            alpha=0.75,
            zorder=1,
        )

# Plot angular polylines with visible breakpoints.
for method in available_methods:
    style = METHOD_STYLE[method]
    y = log_pivot.loc[method, tasks].values.astype(float)

    ax.plot(
        x,
        y,
        color=style["color"],
        linewidth=style["linewidth"],
        alpha=style["alpha"],
        marker="o",
        markersize=4.0 if method != "BioManus" else 5.0,
        markeredgecolor="white",
        markeredgewidth=0.6,
        label=method,
        zorder=style["zorder"],
    )

# Optional: light area under BioManus.
if "BioManus" in available_methods:
    y_bio = log_pivot.loc["BioManus", tasks].values.astype(float)
    ax.fill_between(
        x,
        y_bio,
        y_min,
        color=METHOD_STYLE["BioManus"]["color"],
        alpha=0.055,
        zorder=2,
    )

# Axes and labels.
ax.set_xlim(-0.25, len(tasks) - 0.75)
ax.set_ylim(y_min, y_max)

ax.set_xticks(x)
ax.set_xticklabels([wrap_label(t) for t in tasks], rotation=30, ha="right")

ax.set_yticks([5, 6, 7])
ax.set_yticklabels([r"$10^5$", r"$10^6$", r"$10^7$"])

ax.set_ylabel("Prompt tokens per task (log scale)")
ax.set_title(
    "Context Consumption Across BioAgentBench Tasks",
    pad=12,
    fontweight="bold",
)

# Clean spines.
ax.spines["top"].set_visible(False)
ax.spines["right"].set_visible(False)
ax.spines["left"].set_color("#7A6A5A")
ax.spines["bottom"].set_color("#7A6A5A")

ax.tick_params(axis="x", length=0, pad=8)
ax.tick_params(axis="y", colors="#4B423A")

# Legend.
ax.legend(
    loc="upper center",
    bbox_to_anchor=(0.5, -0.25),
    ncol=6,
    frameon=False,
    handlelength=2.0,
    columnspacing=1.10,
)

fig.tight_layout(rect=[0, 0.10, 1, 1])

# 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("\nTask order from low to high context:")
for i, task in enumerate(tasks, 1):
    print(f"{i:02d}. {task}")

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)}")