#!/usr/bin/env python3
"""Chart measured EvalPlus pass@1 values; label missing suites explicitly."""

from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

from common import read_json


parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--report", type=Path, required=True)
parser.add_argument("--out", type=Path, required=True)
args = parser.parse_args()
benchmarks = read_json(args.report).get("benchmarks") or {}
spec = [("HumanEval", benchmarks.get("humaneval")), ("MBPP", benchmarks.get("mbpp"))]
base = [100 * item["pass_at_1"]["base"] if item else np.nan for _, item in spec]
plus = [100 * item["pass_at_1"]["plus"] if item else np.nan for _, item in spec]
x = np.arange(len(spec))

fig, ax = plt.subplots(figsize=(8, 4.8))
width = 0.34
first = ax.bar(x - width / 2, base, width, label="base tests", color="#2563eb")
second = ax.bar(x + width / 2, plus, width, label="plus tests", color="#059669")
for index, (_, item) in enumerate(spec):
    if item:
        ax.text(first[index].get_x() + first[index].get_width() / 2, base[index] + 1, f"{base[index]:.1f}%", ha="center")
        ax.text(second[index].get_x() + second[index].get_width() / 2, plus[index] + 1, f"{plus[index]:.1f}%", ha="center")
    else:
        ax.text(index, 4, "not measured", rotation=90, ha="center", va="bottom", color="#64748b")
ax.set(title="v25 EvalPlus pass@1", ylabel="pass@1 (%)", xticks=x, xticklabels=[name for name, _ in spec], ylim=(0, 105))
ax.grid(axis="y", alpha=0.25)
ax.legend()
fig.tight_layout()
args.out.parent.mkdir(parents=True, exist_ok=True)
fig.savefig(args.out, dpi=160, bbox_inches="tight")
plt.close(fig)
print(args.out)