#!/usr/bin/env python3
"""Generate the wall-clock / parameters / verbosity plots from models.json + results.json.

Usage
-----
    python scripts/make_plots.py --models models.json --results results.json --out plots

Requires matplotlib. Time plots use GB10 runs only (Thor wall-clock is not comparable).
Each point is the best valid run for that model/quant, matching the table.
"""
from __future__ import annotations

import argparse
import json
import os

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


def best_run(model: dict, runs: list[dict]) -> dict | None:
    matches = [r for r in runs if r["label"].startswith(model["group"]) and r["total_valid"] > 0]
    if not matches:
        return None
    return max(matches, key=lambda r: (r["total_passed"], r["total_valid"], -r["total_void"]))


def build(models: list[dict], runs: list[dict]) -> list[dict]:
    data = []
    for model in models:
        run = best_run(model, runs)
        if not run:
            continue
        data.append({
            **model,
            "box": run["box"],
            "wall_h": run["total_wall_s"] / 3600,
            "tokens": run["total_tokens_out"],
            "tokens_per_task_k": run["total_tokens_out"] / max(1, run["total_valid"]) / 1000,
            "tok_s": run["total_tokens_out"] / max(1, run["total_wall_s"]),
            "score": run["total_passed"] / run["total_valid"],
            "valid": run["total_valid"],
            "void": run["total_void"],
        })
    return data


def scatter(points: list[dict], xkey: str, ykey: str, title: str, xlabel: str, ylabel: str,
            out: str, logx: bool = True) -> None:
    fig, ax = plt.subplots(figsize=(8, 5.5), dpi=150)
    xs = [p[xkey] for p in points]
    ys = [p[ykey] for p in points]
    sc = ax.scatter(xs, ys, c=[p["score"] for p in points], cmap="RdYlGn", vmin=0.9, vmax=1.0,
                    s=75, edgecolors="black", linewidths=0.4, zorder=3)
    if logx:
        ax.set_xscale("log")
    for p in points:
        if p.get("annotate"):
            dx, dy = p.get("offset", (6, 5))
            ax.annotate(p.get("short", p["model"]), (p[xkey], p[ykey]), textcoords="offset points",
                        xytext=(dx, dy), ha=("right" if dx < 0 else "left"), fontsize=7, alpha=0.9)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.grid(True, which="both", alpha=0.25, zorder=0)
    ax.margins(x=0.30)
    cb = fig.colorbar(sc, ax=ax, pad=0.02)
    cb.set_label("valid-task score")
    fig.tight_layout()
    fig.savefig(out, bbox_inches="tight")
    plt.close(fig)


def scatter_country(points: list[dict], xkey: str, ykey: str, title: str, xlabel: str, ylabel: str,
                    out: str, logx: bool = True) -> None:
    from matplotlib.lines import Line2D
    from matplotlib.patches import Patch

    fig, ax = plt.subplots(figsize=(9.5, 6.3), dpi=150)
    colors = {
        "US": "#2E5FA3",
        "China": "#C0392B",
        "Germany": "#E6A817",
        "Czechia": "#7D5BA6",
        "UAE": "#2E8B57",
        "Unknown": "#777777",
        "Abliterated": "#888888",
    }
    def color(country):
        return colors.get(country, colors["Unknown"])
    def base_of(p):
        return p.get("base_country", p.get("country", "Unknown"))
    def post_of(p):
        return p.get("post_country", p.get("country", "Unknown"))
    for p in points:
        ax.scatter(p[xkey], p[ykey], s=90, marker="o",
                   facecolors=color(base_of(p)), edgecolors=color(post_of(p)),
                   linewidths=2.0, zorder=3)
    if logx:
        ax.set_xscale("log")
    for p in points:
        if p.get("annotate"):
            dx, dy = p.get("offset", (6, 5))
            ax.annotate(p.get("short", p["model"]), (p[xkey], p[ykey]), textcoords="offset points",
                        xytext=(dx, dy), ha=("right" if dx < 0 else "left"), fontsize=7, alpha=0.9)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.text(0.01, 0.01, "Wall time includes tasks voided by timeout/truncation. "
            "Fill = pretrained base; edge = post-training/compression.",
            transform=ax.transAxes, fontsize=7, alpha=0.7)
    ax.grid(True, which="both", alpha=0.25, zorder=0)
    ax.margins(x=0.30)
    ax.set_ylim(bottom=0)

    base_present = sorted({base_of(p) for p in points})
    post_present = sorted({post_of(p) for p in points})
    fill_handles = [Patch(facecolor=color(c), edgecolor="black", label=c) for c in base_present]
    edge_handles = [Line2D([0], [0], marker="o", linestyle="none", markerfacecolor="white",
                           markeredgecolor=color(c), markeredgewidth=2.0, markersize=8, label=c)
                    for c in post_present]
    leg1 = ax.legend(handles=fill_handles, title="Base (fill)", fontsize=7, title_fontsize=7,
                     loc="upper left", framealpha=0.95)
    ax.add_artist(leg1)
    ax.legend(handles=edge_handles, title="Post-train (edge)", fontsize=7, title_fontsize=7,
              loc="upper right", framealpha=0.95)
    fig.tight_layout()
    fig.savefig(out, bbox_inches="tight")
    plt.close(fig)


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--models", default="models.json")
    ap.add_argument("--results", default="results.json")
    ap.add_argument("--out", default="plots")
    args = ap.parse_args()

    models = json.load(open(args.models))["models"]
    runs = json.load(open(args.results))["runs"]
    data = build(models, runs)
    gb = [d for d in data if d["box"] == "gb10"]
    os.makedirs(args.out, exist_ok=True)

    scatter(gb, "active_b", "wall_h",
            "Wall clock vs active parameters (GB10, 42 tasks)",
            "Active parameters (B, log)", "Wall clock (hours)",
            os.path.join(args.out, "wall_vs_active.png"))
    perfect = [d for d in gb if d["valid"] > 0 and d["score"] >= 0.999999]
    # Spark-X2.5 is a 14.5 h BF16 outlier and compresses the y-scale for the origin plots.
    origin_pts = [d for d in perfect if not d["model"].startswith("Spark-X2.5")]
    scatter_country(origin_pts, "active_b", "wall_h",
            "Wall clock vs active parameters — perfect-score runs (GB10)",
            "Active parameters (B, log)", "Wall clock (hours)",
            os.path.join(args.out, "wall_vs_active_perfect_country.png"))
    scatter_country(origin_pts, "total_b", "wall_h",
            "Wall clock vs total parameters — perfect-score runs (GB10)",
            "Total parameters (B, log)", "Wall clock (hours)",
            os.path.join(args.out, "wall_vs_total_perfect_country.png"))
    scatter(gb, "total_b", "wall_h",
            "Wall clock vs total parameters (GB10, 42 tasks)",
            "Total parameters (B, log)", "Wall clock (hours)",
            os.path.join(args.out, "wall_vs_total.png"))
    scatter(gb, "tokens", "wall_h",
            "Wall clock vs output tokens (GB10, 42 tasks)",
            "Output tokens (thousands, log)", "Wall clock (hours)",
            os.path.join(args.out, "wall_vs_tokens.png"))

    # verbosity: all boxes, marker = box
    fig, ax = plt.subplots(figsize=(8, 5.5), dpi=150)
    last_sc = None
    for box, marker, label in (("gb10", "o", "GB10"), ("thor", "s", "Thor")):
        pts = [d for d in data if d["box"] == box]
        if not pts:
            continue
        last_sc = ax.scatter([p["active_b"] for p in pts], [p["tokens_per_task_k"] for p in pts],
                             c=[p["score"] for p in pts], cmap="RdYlGn", vmin=0.9, vmax=1.0,
                             marker=marker, s=75, edgecolors="black", linewidths=0.4,
                             label=label, zorder=3)
    for p in data:
        if p.get("annotate"):
            dx, dy = p.get("offset", (6, 5))
            ax.annotate(p.get("short", p["model"]), (p["active_b"], p["tokens_per_task_k"]),
                        textcoords="offset points", xytext=(dx, dy),
                        ha=("right" if dx < 0 else "left"), fontsize=7)
    ax.set_xscale("log")
    ax.set_xlabel("Active parameters (B, log)")
    ax.set_ylabel("Output tokens per valid task (thousands)")
    ax.set_title("Verbosity vs active parameters (42 tasks)")
    ax.grid(True, which="both", alpha=0.25)
    ax.margins(x=0.30)
    ax.legend()
    if last_sc is not None:
        cb = fig.colorbar(last_sc, ax=ax, pad=0.02)
        cb.set_label("valid-task score")
    fig.tight_layout()
    fig.savefig(os.path.join(args.out, "tokens_per_task_vs_active.png"), bbox_inches="tight")
    plt.close(fig)

    # performance per byte: capability and throughput vs quant file size
    sized = [d for d in data if d.get("size_gb")]
    gb_sized = [d for d in sized if d["box"] == "gb10"]
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 5.5), dpi=150)

    for box, marker, label in (("gb10", "o", "GB10"), ("thor", "s", "Thor")):
        pts = [d for d in sized if d["box"] == box]
        if not pts:
            continue
        ax1.scatter([d["size_gb"] for d in pts], [d["score"] * 100 for d in pts],
                    marker=marker, s=75, edgecolors="black", linewidths=0.4, label=label, zorder=3)
    for d in sized:
        if d.get("annotate") and (d["score"] < 0.995 or d.get("highlight")):
            dx, dy = d.get("size_offset", d.get("offset", (6, 5)))
            ax1.annotate(d.get("short", d["model"]), (d["size_gb"], d["score"] * 100),
                         textcoords="offset points", xytext=(dx, dy),
                         ha=("right" if dx < 0 else "left"), fontsize=7)
    ax1.set_xscale("log")
    ax1.set_xlabel("Quant file size (GB, log)")
    ax1.set_ylabel("Valid-task score (%)")
    ax1.set_title("Capability per byte")
    ax1.grid(True, which="both", alpha=0.25)
    ax1.margins(x=0.30)
    ax1.legend()

    sc = ax2.scatter([d["size_gb"] for d in gb_sized], [d["tok_s"] for d in gb_sized],
                     c=[d["score"] for d in gb_sized], cmap="RdYlGn", vmin=0.9, vmax=1.0,
                     s=75, edgecolors="black", linewidths=0.4, zorder=3)
    for d in gb_sized:
        if d.get("annotate"):
            dx, dy = d.get("size_offset", d.get("offset", (6, 5)))
            ax2.annotate(d.get("short", d["model"]), (d["size_gb"], d["tok_s"]),
                         textcoords="offset points", xytext=(dx, dy),
                         ha=("right" if dx < 0 else "left"), fontsize=7)
    ax2.set_xscale("log")
    ax2.set_xlabel("Quant file size (GB, log)")
    ax2.set_ylabel("Output tokens/s (end-to-end)")
    ax2.set_title("Throughput per byte (GB10)")
    ax2.grid(True, which="both", alpha=0.25)
    ax2.margins(x=0.30)
    cb = fig.colorbar(sc, ax=ax2, pad=0.02)
    cb.set_label("valid-task score")

    fig.suptitle("Performance per byte (42 tasks)")
    fig.tight_layout()
    fig.savefig(os.path.join(args.out, "perf_per_byte.png"), bbox_inches="tight")
    plt.close(fig)

    print(f"wrote 7 plots to {args.out} ({len(gb)} GB10 / {len(data) - len(gb)} Thor rows)")


if __name__ == "__main__":
    main()