#!/usr/bin/env python3
"""Render the Beam/Pi pilot from aggregate results, never from invented trajectories."""
from __future__ import annotations

import argparse
import csv
import json
import math
from pathlib import Path


AXES = {
    "reference_seconds": ("reference", "Normalized latency per task (seconds)",
                          "Estimated new-context input; fixed reference rates + measured tool time"),
    "wall_seconds": ("wall", "Elapsed solver time per task (seconds)",
                     "Observed wall time, including API and tool delays"),
    "api_prefill_reference_seconds": ("api-prefill", "API-prefill normalized latency per task (seconds)",
                                      "Sensitivity view: all API prompt tokens + output + measured tool time"),
}
CONDITIONS = ("single", "peers", "async")
STYLES = {
    "single": ("#f66535", "solid", "-", "Single agent · 1M estimated context tokens"),
    "peers": ("#c65329", "dot", ":", "Five-agent team · 200k estimated context tokens per agent"),
    "async": ("#fb9a7d", "dashdot", "-.", "Async subagents · 200k estimated context tokens per agent"),
}
TITLE = "Multi-agent ProgramBench: score vs. latency"
MODEL = "Reflection Beam-501B-A23B + Pi"
SOURCE = "https://www-cdn.anthropic.com/fc1b44717c85dc068bc6ba5024219938094694bd/Claude%20Opus%205.5%20System%20Card.pdf"


def numeric(value):
    try:
        result = float(value)
    except (ValueError, TypeError):
        raise ValueError("Expected finite numeric curve data") from None
    if isinstance(value, bool) or not math.isfinite(result):
        raise ValueError("Expected finite numeric curve data")
    return result


def is_true(value):
    return value is True or value == "True"


def count(value):
    value = numeric(value)
    if value != int(value) or not 0 <= value <= 5:
        raise ValueError("Task coverage must be an integer from zero to five")
    return int(value)


def figure_data(summary, curves, axis="reference_seconds"):
    if axis not in AXES:
        raise ValueError("Unknown latency axis")
    summary = summary or {}
    groups = {}
    for condition in CONDITIONS:
        episodes = [e for e in summary.get("episodes", []) if e.get("condition") == condition]
        observed = len({e.get("task") for e in episodes if e.get("has_valid_grade") is True})
        invalid = sum(int(e.get("invalid_grade_count", 0)) for e in episodes)
        finished = len(episodes) == 5 and all(e.get("status") in {"completed", "budget_exhausted"} for e in episodes)
        rows = []
        clock_ok = axis == "wall_seconds" or (bool(episodes) and all(e.get("clock_valid") is True for e in episodes))
        for row in curves or []:
            if row.get("axis") != axis or row.get("condition") != condition:
                continue
            if count(row["tasks_total"]) != 5:
                raise ValueError("The pilot plot requires the fixed five-task denominator")
            seconds = numeric(row["time_seconds"])
            coverage = count(row["tasks_with_valid_grade"])
            if seconds < 0:
                raise ValueError("Latency cannot be negative")
            if not is_true(row.get("clock_valid")):
                clock_ok = False
            score = row.get("mean_hidden_test_fraction")
            if score in (None, ""):
                continue
            score = numeric(score)
            if not 0 <= score <= 1:
                raise ValueError("Score outside zero to one")
            # A logarithmic axis has no zero. Never assign an invented positive timestamp.
            if seconds > 0:
                rows.append({"seconds": seconds, "score": score, "coverage": coverage})
        rows.sort(key=lambda row: row["seconds"])
        checkpoints = {numeric(point[axis]) for episode in episodes for point in episode.get("points", [])
                       if axis in point}
        for row in rows:
            row["checkpoint"] = row["seconds"] in checkpoints or len(rows) == 1
        available = observed > 0 and clock_ok and bool(rows)
        groups[condition] = {"rows": rows if available else [], "observed": observed,
                             "invalid_grades": invalid, "finished": finished, "clock_valid": clock_ok,
                             "available": available, "complete": observed == 5 and invalid == 0 and finished}
    values = [row["seconds"] for group in groups.values() for row in group["rows"]]
    limits = (min(values), max(values)) if values else (30.0, 7200.0)
    if limits[0] == limits[1]:
        limits = (limits[0] / 1.1, limits[1] * 1.1)
    has_grades = any(group["observed"] for group in groups.values())
    available = any(group["available"] for group in groups.values())
    if not has_grades:
        status = "Awaiting graded benchmark results"
        detail = "Reference calibration and infrastructure checks are not model benchmark scores."
    elif not available:
        status = "This latency view is unavailable"
        detail = "Complete clock records and positive-time checkpoints are required. Check the wall-time view."
    elif not all(group["clock_valid"] for group in groups.values()):
        status = "One or more normalized curves are unavailable"
        detail = "Missing or invalid clock records suppress the affected condition; check the wall-time view."
    elif not all(group["complete"] for group in groups.values()):
        status = "Provisional results · check grading coverage"
        detail = "Five-task denominator; zero before any valid grade; last valid score retained after grading failure."
    else:
        status = "All 15 episodes graded"
        detail = "Equal task weights; latest valid score carried forward; regressions retained."
    rates = summary.get("reference_clock", {})
    prefill = numeric(rates.get("prefill_tokens_per_second", 10000))
    decode = numeric(rates.get("decode_tokens_per_second", 265))
    if prefill <= 0 or decode <= 0:
        raise ValueError("Reference rates must be positive")
    return {"axis": axis, "groups": groups, "limits": limits, "available": available,
            "has_grades": has_grades, "status": status, "detail": detail,
            "rates": f"Reference rates: {prefill:,.0f} input and {decode:,.0f} output tokens/s (our assumptions)."}


def threshold_comparison(data, threshold=0.6):
    """First observed crossing, never interpolation/extrapolation or a copied paper ratio."""
    crossing = {}
    for condition in ("single", "peers"):
        group = data["groups"][condition]
        if not group["complete"] or not group["available"]:
            return None
        hit = next((row for row in group["rows"] if row["score"] >= threshold and row["coverage"] == 5), None)
        if hit is None:
            return None
        crossing[condition] = hit["seconds"]
    ratio = crossing["single"] / crossing["peers"]
    comparison = f"{ratio:.2f}× faster" if ratio > 1.005 else f"{1 / ratio:.2f}× slower" if ratio < 0.995 else "the same observed latency"
    return {"threshold": threshold, **crossing, "ratio": ratio,
            "label": f"Time to {threshold:g}: five-agent team {comparison} (first graded crossing)"}


def notes(data):
    result = ["5 selected tasks · 1 repetition · No confidence intervals · 8M API-token cap per episode",
              data["detail"], "Coverage counts tasks with a valid grade by each cutoff; it is not the number still running."]
    if data["axis"] != "wall_seconds":
        result.extend([data["rates"], "Only completed operations enter the clock. This is not Anthropic's exact latency calibration."])
    return result


def legend_suffix(group):
    return " · pending" if not group["observed"] else " · clock unavailable" if not group["clock_valid"] else ""


def ticks(limits):
    lo, hi = limits
    values = [multiple * 10.0 ** power for power in range(-3, 9) for multiple in (1, 3)
              if lo <= multiple * 10.0 ** power <= hi]
    return values, [f"{x:,.0f}" if x >= 1 else f"{x:g}" for x in values]


def build_plotly(summary, curves, axis="reference_seconds"):
    data = figure_data(summary, curves, axis)
    traces = []
    for condition in CONDITIONS:
        group = data["groups"][condition]
        color, dash, _, label = STYLES[condition]
        rows = group["rows"]
        suffix = legend_suffix(group)
        traces.append({"type": "scatter", "mode": "lines+markers", "name": label + suffix,
                       "x": [r["seconds"] for r in rows] or [None], "y": [r["score"] for r in rows] or [None],
                       "customdata": [[r["coverage"]] for r in rows],
                       "line": {"color": color, "width": 3, "dash": dash, "shape": "hv"},
                       "marker": {"color": color, "size": [5 if r["checkpoint"] else 0 for r in rows] or [0]},
                       "connectgaps": False, "legendgroup": condition,
                       "hovertemplate": "%{x:,.1f}s · mean score %{y:.3f}<br>%{customdata[0]}/5 tasks graded<extra>%{fullData.name}</extra>"})
        if rows:
            traces.append({"type": "scatter", "mode": "lines+markers", "x": [r["seconds"] for r in rows],
                           "y": [r["coverage"] for r in rows], "yaxis": "y2", "showlegend": False,
                           "legendgroup": condition, "name": label,
                           "line": {"color": color, "width": 2, "dash": dash, "shape": "hv"},
                           "marker": {"color": color, "size": [3 if r["checkpoint"] else 0 for r in rows]},
                           "hovertemplate": "%{x:,.1f}s · %{y}/5 tasks graded<extra></extra>"})
    annotations = [{"text": data["status"], "xref": "paper", "yref": "paper", "x": 0, "y": 1.025,
                    "showarrow": False, "xanchor": "left", "yanchor": "bottom", "font": {"size": 14, "color": "#61615d"}}]
    if not data["available"]:
        annotations.append({"text": "No benchmark curve to display yet" if not data["has_grades"] else "No valid curve for this clock",
                            "xref": "paper", "yref": "paper", "x": 0.5, "y": 0.64, "showarrow": False,
                            "font": {"size": 21, "color": "#767670"}})
    crossing = threshold_comparison(data)
    if crossing:
        annotations.extend([{"x": math.log10(crossing["single"]), "ax": math.log10(crossing["peers"]), "y": 0.6, "ay": 0.6,
                             "xref": "x", "axref": "x", "yref": "y", "ayref": "y", "arrowside": "end+start",
                             "arrowhead": 2, "startarrowhead": 2, "arrowcolor": "#777", "text": ""},
                            {"text": crossing["label"], "xref": "paper", "yref": "paper", "x": 0.02, "y": 0.78,
                             "showarrow": False, "xanchor": "left", "bgcolor": "rgba(255,255,255,0.92)"}])
    tick_values, tick_text = ticks(data["limits"])
    annotations.append({"text": "<br>".join(notes(data)), "xref": "paper", "yref": "paper", "x": 0, "y": -0.22,
                        "showarrow": False, "xanchor": "left", "yanchor": "top", "align": "left",
                        "font": {"size": 11, "color": "#666660"}})
    return {"data": traces, "layout": {
        "title": {"text": TITLE + "<br><sup>" + MODEL + " · " + AXES[axis][2] + "</sup>", "x": 0.07,
                  "y": 0.95, "yanchor": "top",
                  "font": {"size": 24, "color": "#1b1b18"}},
        "height": 850, "paper_bgcolor": "white", "plot_bgcolor": "white",
        "font": {"family": "Arial, sans-serif", "color": "#22221f"},
        "margin": {"t": 200, "b": 200, "l": 85, "r": 25},
        "legend": {"orientation": "v", "x": 0, "y": 1.275, "yanchor": "top", "font": {"size": 12}},
        "xaxis": {"type": "log", "range": [math.log10(data["limits"][0]), math.log10(data["limits"][1])],
                  "title": {"text": AXES[axis][1]}, "anchor": "y2", "tickmode": "array", "tickvals": tick_values,
                  "ticktext": tick_text, "showgrid": False, "showline": True, "linecolor": "#aaa9a2"},
        "yaxis": {"domain": [0.27, 1], "range": [0, 1], "dtick": 0.2, "tickformat": ".1f",
                  "title": {"text": "Fraction of hidden tests passed"}, "gridcolor": "#e1e0da", "zeroline": False},
        "yaxis2": {"domain": [0, 0.12], "range": [-0.15, 5.2], "tickvals": [0, 5],
                   "title": {"text": "Tasks graded", "font": {"size": 11}}, "gridcolor": "#efeee8", "zeroline": False},
        "annotations": annotations, "hovermode": "closest",
        "meta": {"status": data["status"], "has_results": data["available"], "axis": axis,
                 "notes": notes(data), "threshold_comparison": crossing, "source": SOURCE,
                 "coverage": {c: {k: v for k, v in g.items() if k != "rows"} for c, g in data["groups"].items()}}
    }}


def render_outputs(summary, curves, output_dir):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from matplotlib.lines import Line2D
    from matplotlib.ticker import FixedLocator, FuncFormatter, NullFormatter
    directory = Path(output_dir)
    directory.mkdir(parents=True, exist_ok=True)
    outputs = {}
    with plt.rc_context({"font.family": "DejaVu Sans", "font.size": 11, "svg.fonttype": "none"}):
        for axis, (slug, xlabel, subtitle) in AXES.items():
            data = figure_data(summary, curves, axis)
            fig = plt.figure(figsize=(13.5, 8.7), facecolor="white")
            ax = fig.add_axes((0.09, 0.33, 0.875, 0.435))
            coverage = fig.add_axes((0.09, 0.22, 0.875, 0.065), sharex=ax)
            for chart in (ax, coverage):
                chart.set_xscale("log")
                chart.set_xlim(*data["limits"])
                chart.grid(axis="y", color="#e1e0da", linewidth=0.8)
                chart.spines[["top", "right"]].set_visible(False)
                chart.spines[["left", "bottom"]].set_color("#aaa9a2")
                chart.tick_params(colors="#393933", labelsize=10)
            ax.set_ylim(0, 1)
            ax.set_yticks([i / 5 for i in range(6)])
            ax.set_ylabel("Fraction of hidden tests passed", labelpad=10)
            ax.tick_params(axis="x", which="both", labelbottom=False)
            coverage.set_ylim(-0.15, 5.2)
            coverage.set_yticks([0, 5])
            coverage.set_ylabel("Tasks\ngraded", fontsize=9, labelpad=17)
            coverage.set_xlabel(xlabel, labelpad=9)
            tick_values, _ = ticks(data["limits"])
            coverage.xaxis.set_major_locator(FixedLocator(tick_values))
            coverage.xaxis.set_major_formatter(FuncFormatter(lambda value, _: f"{value:,.0f}" if value >= 1 else f"{value:g}"))
            coverage.xaxis.set_minor_formatter(NullFormatter())
            handles = []
            for condition in CONDITIONS:
                color, _, linestyle, label = STYLES[condition]
                handles.append(Line2D([0], [0], color=color, linestyle=linestyle, linewidth=2.6,
                                      label=label + legend_suffix(data["groups"][condition])))
                rows = data["groups"][condition]["rows"]
                if rows:
                    xs = [r["seconds"] for r in rows]
                    ax.step(xs, [r["score"] for r in rows], where="post", color=color, linestyle=linestyle, linewidth=2.6)
                    coverage.step(xs, [r["coverage"] for r in rows], where="post", color=color, linestyle=linestyle, linewidth=1.8)
                    marked = [r for r in rows if r["checkpoint"]]
                    ax.scatter([r["seconds"] for r in marked], [r["score"] for r in marked], color=color, s=13, zorder=3)
                    coverage.scatter([r["seconds"] for r in marked], [r["coverage"] for r in marked], color=color, s=7, zorder=3)
            fig.text(0.09, 0.94, TITLE, fontsize=24, fontweight="bold", color="#1b1b18")
            fig.text(0.09, 0.901, MODEL + " · " + subtitle, fontsize=11, color="#55554f")
            fig.legend(handles=handles, loc="upper left", bbox_to_anchor=(0.085, 0.884), frameon=False, fontsize=10)
            fig.text(0.09, 0.787, data["status"], fontsize=10, color="#686861")
            if not data["available"]:
                ax.text(0.5, 0.53, "Awaiting graded benchmark results" if not data["has_grades"] else "No valid curve for this clock",
                        transform=ax.transAxes, ha="center", color="#797971", fontsize=17)
            crossing = threshold_comparison(data)
            if crossing:
                ax.annotate("", xy=(crossing["single"], 0.6), xytext=(crossing["peers"], 0.6),
                            arrowprops={"arrowstyle": "<->", "color": "#777"})
                ax.text(0.02, 0.7, crossing["label"], transform=ax.transAxes, fontsize=10,
                        bbox={"facecolor": "white", "edgecolor": "none", "alpha": 0.92})
            for index, note in enumerate(notes(data)):
                fig.text(0.09, 0.127 - index * 0.022, note, fontsize=9, color="#696962")
            prefix = f"score-vs-latency-{axis}"
            for extension in ("png", "svg", "pdf"):
                path = directory / f"{prefix}.{extension}"
                fig.savefig(path, dpi=180, facecolor="white")
                outputs[path.name] = str(path.resolve())
            plt.close(fig)
            path = directory / f"{prefix}.json"
            path.write_text(json.dumps(build_plotly(summary, curves, axis), indent=2, allow_nan=False) + "\n")
            outputs[path.name] = str(path.resolve())
    return outputs


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--experiment-dir", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    args = parser.parse_args()
    summary_path, curves_path = args.experiment_dir / "summary.json", args.experiment_dir / "curves.csv"
    summary = json.loads(summary_path.read_text()) if summary_path.exists() else None
    curves = list(csv.DictReader(curves_path.open())) if curves_path.exists() else []
    print(json.dumps(render_outputs(summary, curves, args.output_dir), indent=2))


if __name__ == "__main__":
    main()