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

import csv
import math
from pathlib import Path

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


ROOT = Path(__file__).resolve().parent
CSV_PATH = ROOT / "biomni_context_tokens.csv"
OUT_PREFIX = ROOT / "biomni_context_vs_tool_scale_clean"
STATS_PATH = ROOT / "biomni_context_vs_tool_scale_stats.csv"

SCALE_ORDER = [0, 100, 500, 1000, 2000]
SCALE_LABELS = {
    0: "No MCP",
    100: "100",
    500: "500",
    1000: "1k",
    2000: "2k",
}


def clean(value: str | None) -> str:
    return (value or "").replace("\ufeff", "").strip()


def load_rows(path: Path) -> list[dict]:
    rows: list[dict] = []
    current_scale: int | None = None

    with path.open("r", encoding="utf-8-sig", newline="") as handle:
        reader = csv.reader(handle)
        for raw in reader:
            if not raw:
                continue

            first = clean(raw[0])
            second = clean(raw[1] if len(raw) > 1 else "")
            if first.startswith("Experiments") or second == "Tasks":
                continue

            if first:
                try:
                    current_scale = int(float(first))
                except ValueError:
                    continue

            if current_scale is None or len(raw) < 6:
                continue

            task = clean(raw[1])
            if not task:
                continue

            try:
                prompt_tokens = float(clean(raw[3]).replace(",", ""))
                completion_tokens = float(clean(raw[4]).replace(",", ""))
                total_tokens = float(clean(raw[5]).replace(",", ""))
            except ValueError:
                continue

            rows.append(
                {
                    "scale": current_scale,
                    "task": task,
                    "results_match": clean(raw[2]).upper() == "TRUE",
                    "prompt_tokens": prompt_tokens,
                    "completion_tokens": completion_tokens,
                    "total_tokens": total_tokens,
                }
            )
    return rows


def geometric_mean(values: list[float]) -> float:
    values = [v for v in values if v > 0]
    return float(math.exp(sum(math.log(v) for v in values) / len(values))) if values else 0.0


def token_formatter(value: float, _pos: int) -> 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 main() -> int:
    rows = load_rows(CSV_PATH)

    by_scale = {scale: [] for scale in SCALE_ORDER}
    by_task: dict[str, dict[int, float]] = {}
    for row in rows:
        scale = row["scale"]
        if scale not in by_scale:
            continue
        value = row["prompt_tokens"]
        by_scale[scale].append(value)
        by_task.setdefault(row["task"], {})[scale] = value

    scales = [scale for scale in SCALE_ORDER if by_scale.get(scale)]
    data = [by_scale[scale] for scale in scales]
    positions = np.arange(len(scales), dtype=float)

    medians = np.array([np.median(values) for values in data])
    means = np.array([np.mean(values) for values in data])
    geo_means = np.array([geometric_mean(values) for values in data])
    q1 = np.array([np.percentile(values, 25) for values in data])
    q3 = np.array([np.percentile(values, 75) for values in data])

    # Clean conference-style figure settings.
    plt.rcParams.update(
        {
            "font.family": "DejaVu Sans",
            "font.size": 10.5,
            "axes.labelsize": 11,
            "axes.titlesize": 12,
            "xtick.labelsize": 10,
            "ytick.labelsize": 10,
            "legend.fontsize": 9.5,
            "axes.spines.top": False,
            "axes.spines.right": False,
            "pdf.fonttype": 42,
            "ps.fonttype": 42,
        }
    )

    fig, (ax_dist, ax_trend) = plt.subplots(
        1,
        2,
        figsize=(9.2, 4.8),
        gridspec_kw={"width_ratios": [1.25, 1.0]},
        constrained_layout=True,
    )

    box_color = "#E6EEF5"
    edge_color = "#243447"
    point_color = "#1F2933"
    median_color = "#E4572E"
    mean_color = "#0B1F33"
    geomean_color = "#0B1F33"
    iqr_color = "#2E86AB"
    traj_color = "#A7B0BD"

    # ------------------------------------------------------------------
    # Left panel: boxplot + jittered raw task points + geometric mean.
    # Violin plots are removed because each scale has only 10 tasks.
    # ------------------------------------------------------------------
    box = ax_dist.boxplot(
        data,
        positions=positions,
        widths=0.46,
        patch_artist=True,
        showfliers=False,
        medianprops={"color": median_color, "linewidth": 1.8},
        whiskerprops={"color": edge_color, "linewidth": 1.0},
        capprops={"color": edge_color, "linewidth": 1.0},
        boxprops={"edgecolor": edge_color, "linewidth": 1.0},
    )
    for patch in box["boxes"]:
        patch.set_facecolor(box_color)
        patch.set_alpha(0.95)

    rng = np.random.default_rng(20260523)
    for i, values in enumerate(data):
        jitter = rng.normal(positions[i], 0.055, size=len(values))
        ax_dist.scatter(
            jitter,
            values,
            s=24,
            color=point_color,
            alpha=0.70,
            linewidth=0.35,
            edgecolor="white",
            zorder=4,
        )

    ax_dist.plot(
        positions,
        geo_means,
        color=geomean_color,
        linewidth=2.0,
        marker="D",
        markersize=5.2,
        label="Geometric mean",
        zorder=5,
    )
    ax_dist.set_yscale("log")
    ax_dist.yaxis.set_major_formatter(FuncFormatter(token_formatter))
    ax_dist.set_xticks(positions)
    ax_dist.set_xticklabels([SCALE_LABELS[scale] for scale in scales])
    ax_dist.set_xlabel("Available MCP tools")
    ax_dist.set_ylabel("Prompt tokens per task")
    ax_dist.set_title("Task-level distribution")
    ax_dist.grid(axis="y", which="major", linestyle="-", linewidth=0.55, alpha=0.25)
    ax_dist.grid(axis="y", which="minor", linestyle=":", linewidth=0.4, alpha=0.18)
    ax_dist.legend(frameon=False, loc="upper left")

    # ------------------------------------------------------------------
    # Right panel: aggregate trend. Keep the panel simple; no annotation box.
    # ------------------------------------------------------------------
    for task, values_by_scale in sorted(by_task.items()):
        y = [values_by_scale.get(scale, np.nan) for scale in scales]
        if np.isnan(y).any():
            continue
        ax_trend.plot(positions, y, color=traj_color, alpha=0.22, linewidth=0.9, zorder=1)

    ax_trend.fill_between(positions, q1, q3, color=iqr_color, alpha=0.15, label="IQR", zorder=2)
    ax_trend.plot(
        positions,
        medians,
        color=median_color,
        linewidth=2.2,
        marker="o",
        markersize=5.2,
        label="Median",
        zorder=4,
    )
    ax_trend.plot(
        positions,
        means,
        color=mean_color,
        linewidth=1.8,
        marker="s",
        markersize=4.8,
        linestyle="--",
        label="Mean",
        zorder=4,
    )

    ax_trend.set_yscale("log")
    ax_trend.yaxis.set_major_formatter(FuncFormatter(token_formatter))
    ax_trend.set_xticks(positions)
    ax_trend.set_xticklabels([SCALE_LABELS[scale] for scale in scales])
    ax_trend.set_xlabel("Available MCP tools")
    ax_trend.set_title("Aggregate trend")
    ax_trend.grid(axis="y", which="major", linestyle="-", linewidth=0.55, alpha=0.25)
    ax_trend.grid(axis="y", which="minor", linestyle=":", linewidth=0.4, alpha=0.18)
    ax_trend.legend(frameon=False, loc="upper left")

    # Short title only. Put detailed explanation in the paper caption.
    baseline = geo_means[0]
    final = geo_means[-1]
    fold = final / baseline if baseline else float("nan")
    fig.suptitle(
        f"Biomni Context Consumption vs. MCP Tool Scale ({fold:.1f}x geometric mean)",
        fontsize=13.5,
        fontweight="bold",
    )

    for ext in ("svg", "pdf", "png"):
        fig.savefig(f"{OUT_PREFIX}.{ext}", dpi=360, bbox_inches="tight")

    with STATS_PATH.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.writer(handle)
        writer.writerow(["scale", "label", "mean", "median", "geometric_mean", "q1", "q3", "n_tasks"])
        for scale, values, mean, median, geomean, lo, hi in zip(scales, data, means, medians, geo_means, q1, q3):
            writer.writerow([
                scale,
                SCALE_LABELS[scale],
                f"{mean:.6f}",
                f"{median:.6f}",
                f"{geomean:.6f}",
                f"{lo:.6f}",
                f"{hi:.6f}",
                len(values),
            ])

    print("Scale\tMean\tMedian\tGeomean\tQ1\tQ3\tN")
    for scale, values, mean, median, geomean, lo, hi in zip(scales, data, means, medians, geo_means, q1, q3):
        print(f"{SCALE_LABELS[scale]}\t{mean:.0f}\t{median:.0f}\t{geomean:.0f}\t{lo:.0f}\t{hi:.0f}\t{len(values)}")
    print(f"Saved: {OUT_PREFIX}.svg")
    print(f"Saved: {OUT_PREFIX}.pdf")
    print(f"Saved: {OUT_PREFIX}.png")
    print(f"Saved: {STATS_PATH}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())