"""physics-llm-benchmarks — a tiny companion Space that renders the interactive
benchmark charts for the Physics-LLM blog post. No model, no GPU: it just loads
the precomputed results.json (real LFM2-350M rollouts scored against Pymunk) and
draws three interactive Plotly figures. Embedded in the blog via <gradio-app>.
"""
import json
import os

import gradio as gr
import plotly.graph_objects as go

HERE = os.path.dirname(__file__)

BG = "#0b0f17"
GRID = "#222b3a"
FG = "#c9d1d9"
TRAINED = "#4ea1ff"
HELDOUT = "#ff7c5b"
PALETTE = ["#4ea1ff", "#ff7c5b", "#ffd166", "#06d6a0", "#c77dff", "#ff5dac",
           "#7ee787", "#f78166", "#79c0ff", "#d2a8ff", "#a5d6ff", "#ffab70"]

with open(os.path.join(HERE, "results.json")) as f:
    DATA = json.load(f)


def _layout(fig, title, xt, yt, h=520):
    fig.update_layout(
        title=dict(text=title, font=dict(color=FG, size=17)),
        paper_bgcolor=BG, plot_bgcolor=BG, font=dict(color=FG, size=13),
        height=h, margin=dict(l=70, r=30, t=60, b=60),
        legend=dict(bgcolor="rgba(0,0,0,0)", font=dict(size=11)),
        xaxis=dict(title=xt, gridcolor=GRID, zerolinecolor=GRID),
        yaxis=dict(title=yt, gridcolor=GRID, zerolinecolor=GRID),
    )
    return fig


def fig_error_growth():
    fig = go.Figure()
    items = sorted(DATA["scenarios"].items(),
                   key=lambda kv: (kv[1]["meta"]["held_out"], kv[0]))
    for i, (name, sc) in enumerate(items):
        ho = sc["meta"]["held_out"]
        xs, ys = [], []
        for k, p in enumerate(sc["per_frame"]):
            if p["mean_pct_diag"] is not None:
                xs.append(k + 1); ys.append(p["mean_pct_diag"])
        fig.add_trace(go.Scatter(
            x=xs, y=ys, mode="lines", name=name + (" (held-out)" if ho else ""),
            line=dict(color=PALETTE[i % len(PALETTE)], width=2.2,
                      dash="dash" if ho else "solid"),
            hovertemplate=f"{name}<br>frame %{{x}}<br>%{{y:.2f}}%% diag<extra></extra>"))
    return _layout(fig, "Position drift vs Pymunk ground truth",
                   "rollout step", "mean position error (% of scene diagonal)")


def fig_per_scenario():
    items = sorted(DATA["scenarios"].items(),
                   key=lambda kv: kv[1]["meta"]["mean_dist_pct_diag"] or 0)
    names = [k for k, _ in items]
    vals = [v["meta"]["mean_dist_pct_diag"] for _, v in items]
    cols = [HELDOUT if v["meta"]["held_out"] else TRAINED for _, v in items]
    fig = go.Figure(go.Bar(x=names, y=vals, marker_color=cols,
                           text=[f"{x:.2f}%" for x in vals], textposition="outside",
                           hovertemplate="%{x}<br>%{y:.3f}%% diag<extra></extra>"))
    return _layout(fig, "Mean rollout error per scenario  (blue = trained, orange = held-out)",
                   "scenario", "mean position error (% of scene diagonal)")


def fig_throughput():
    items = list(DATA["scenarios"].items())
    xs = [v["meta"]["n_obj"] for _, v in items]
    ys = [v["meta"]["frame_per_s"] for _, v in items]
    cols = [HELDOUT if v["meta"]["held_out"] else TRAINED for _, v in items]
    names = [k for k, _ in items]
    fig = go.Figure(go.Scatter(
        x=xs, y=ys, mode="markers+text", text=names, textposition="top center",
        textfont=dict(size=9, color=FG),
        marker=dict(size=14, color=cols, line=dict(width=1, color="#0b0f17")),
        hovertemplate="%{text}<br>%{x} objects<br>%{y:.2f} frame/s<extra></extra>"))
    return _layout(fig, "End-to-end rollout rate vs scene complexity  (sglang, bf16)",
                   "objects in scene", "frames / second (incl. prompt rebuild)")


CFG = DATA.get("config", {})
INTRO = f"""### Physics-LLM — interactive benchmarks

Real rollouts of **LFM2-350M** ({CFG.get('precision','bf16')}, served with **{CFG.get('engine','sglang')}**)
scored against the **Pymunk** engine it was distilled from. {len(DATA['scenarios'])} scenarios,
{CFG.get('n_frames','?')} frames each, greedy decoding. Error is the mean object-position distance to
ground truth, as a percentage of the scene diagonal (scale-free). Hover for per-point values.
"""

with gr.Blocks(title="Physics-LLM benchmarks", fill_width=True) as demo:
    gr.Markdown(INTRO)
    with gr.Tab("Drift over time"):
        gr.Plot(fig_error_growth())
    with gr.Tab("Error per scenario"):
        gr.Plot(fig_per_scenario())
    with gr.Tab("Throughput"):
        gr.Plot(fig_throughput())

if __name__ == "__main__":
    demo.launch()