# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License").
"""Build an HTML visual-review report for a completed LIBERO smoke-test run.

Reads the artifacts produced by ``run_10_smoke_tests.py`` (summary.json,
per-scenario metadata.json / rollout_summary.json / actions.npy / video.mp4 /
frames/), computes a few simple action statistics, renders one card per
scenario (with a video preview, links to the raw files, and 4 small plots),
and writes:

    <run-dir>/visual_report.html
    <run-dir>/visual_summary.csv
    <run-dir>/visual_summary.json
    <run-dir>/plots/<scenario_id>/{action_norm,action_mean_per_dof,gripper_over_time,action_delta_norm}.png

It does NOT touch any existing rollout output, does NOT need a GPU, and does
NOT rerun LIBERO. Missing files become a warning inside the relevant card
instead of crashing.

CLI::

    python examples/LIBERO/smoke_tests/visualize_smoke_run.py \
        --run-dir outputs/libero_smoke_tests/20260512_122756 --open

Flags: --run-dir  --open  --no-plots  --no-video-embed
"""

from __future__ import annotations

import argparse
import base64
import csv
import datetime as _dt
import html
import json
import os
from pathlib import Path
import sys
import tempfile
import traceback

# matplotlib must be configured before import; use a headless backend + a
# writable config dir (the default ~/.config/matplotlib may be read-only).
os.environ.setdefault("MPLCONFIGDIR", tempfile.mkdtemp(prefix="mpl-visualize-"))
import numpy as np

_MPL_OK = True
_MPL_ERR = None
try:
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
except Exception as e:  # noqa: BLE001
    _MPL_OK = False
    _MPL_ERR = repr(e)

# The LIBERO action dict has 7 keys; `_libero_rollout_worker._action_to_numpy`
# concatenates them in *sorted (alphabetical)* key order, so the columns of
# actions.npy for this run are exactly:
LIBERO_DOF_NAMES = [
    "action.gripper",  # col 0
    "action.pitch",    # col 1
    "action.roll",     # col 2
    "action.x",        # col 3
    "action.y",        # col 4
    "action.yaw",      # col 5
    "action.z",        # col 6
]
LIBERO_GRIPPER_COL = 0


# --------------------------------------------------------------------------- #
# small IO helpers
# --------------------------------------------------------------------------- #
def _load_json(path: Path):
    try:
        with open(path) as f:
            return json.load(f), None
    except FileNotFoundError:
        return None, f"missing file: {path.name}"
    except Exception as e:  # noqa: BLE001
        return None, f"failed to read {path.name}: {e!r}"


def _rel(path: Path, start: Path) -> str:
    try:
        return os.path.relpath(path, start)
    except ValueError:
        return str(path)


def _fmt(v) -> str:
    if v is None:
        return "—"
    if isinstance(v, float):
        return f"{v:.4g}"
    return str(v)


def _b64_data_uri(path: Path, mime: str) -> str | None:
    try:
        data = path.read_bytes()
    except OSError:
        return None
    return f"data:{mime};base64," + base64.b64encode(data).decode("ascii")


# --------------------------------------------------------------------------- #
# action statistics
# --------------------------------------------------------------------------- #
def _squeeze_actions(arr: np.ndarray) -> np.ndarray:
    """Best-effort normalize an actions.npy array to shape (n_calls, horizon, n_dof).

    The worker saves shape (n_calls, batch=1, horizon, n_dof). We also tolerate
    (n_calls, horizon, n_dof) and (n_calls, n_dof).
    """
    a = np.asarray(arr)
    if a.dtype == object:
        # list of heterogeneous things saved with allow_pickle; try to stack
        try:
            a = np.stack([np.asarray(x) for x in a], axis=0)
        except Exception:
            raise ValueError(f"cannot interpret object-dtype actions array of shape {arr.shape}")
    if a.ndim == 4 and a.shape[1] == 1:
        a = a[:, 0, :, :]
    if a.ndim == 2:  # (n_calls, n_dof) -> add a length-1 horizon
        a = a[:, None, :]
    if a.ndim != 3:
        raise ValueError(f"unexpected actions array shape {arr.shape} (squeezed to {a.shape})")
    return a.astype(np.float64)


def _detect_gripper_col(acts3: np.ndarray) -> tuple[int, list[str]]:
    """Return (gripper_col_index, dof_names)."""
    n_dof = acts3.shape[-1]
    if n_dof == 7:
        return LIBERO_GRIPPER_COL, list(LIBERO_DOF_NAMES)
    names = [f"dof_{i}" for i in range(n_dof)]
    # generic heuristic: the DoF whose values most often saturate near +/-1,
    # falling back to the last column ("7th-style" gripper convention).
    flat = acts3.reshape(-1, n_dof)
    sat = (np.abs(flat) > 0.5).mean(axis=0)
    col = int(np.argmax(sat)) if sat.size and sat.max() > 0.5 else n_dof - 1
    return col, names


def compute_action_stats(acts3: np.ndarray) -> dict:
    """acts3: (n_calls, horizon, n_dof) -> dict of arrays + scalars."""
    n_calls, horizon, n_dof = acts3.shape
    gripper_col, dof_names = _detect_gripper_col(acts3)

    # per-policy-call norms
    step0_norm = np.linalg.norm(acts3[:, 0, :], axis=-1)             # (n_calls,) norm of the first executed step
    chunk_mean = acts3.mean(axis=1)                                   # (n_calls, n_dof) mean over the chunk
    chunk_mean_norm = np.linalg.norm(chunk_mean, axis=-1)             # (n_calls,)
    per_step_norm = np.linalg.norm(acts3.reshape(-1, n_dof), axis=-1)  # (n_calls*horizon,)

    # mean (and std) action value per DoF over all calls x steps
    flat = acts3.reshape(-1, n_dof)
    mean_per_dof = flat.mean(axis=0)
    std_per_dof = flat.std(axis=0)

    # gripper command over time: value at chunk-step 0 per call
    gripper_step0 = acts3[:, 0, gripper_col]                          # (n_calls,)
    gripper_all = flat[:, gripper_col]

    # action-delta norm between consecutive *chunks* (Frobenius over horizon x dof)
    if n_calls >= 2:
        delta = acts3[1:] - acts3[:-1]                                # (n_calls-1, horizon, n_dof)
        delta_norm = np.linalg.norm(delta.reshape(n_calls - 1, -1), axis=-1)
    else:
        delta_norm = np.zeros((0,), dtype=np.float64)

    return {
        "n_calls": n_calls,
        "horizon": horizon,
        "n_dof": n_dof,
        "dof_names": dof_names,
        "gripper_col": gripper_col,
        "step0_norm": step0_norm,
        "chunk_mean_norm": chunk_mean_norm,
        "per_step_norm": per_step_norm,
        "mean_per_dof": mean_per_dof,
        "std_per_dof": std_per_dof,
        "gripper_step0": gripper_step0,
        "gripper_all": gripper_all,
        "delta_norm": delta_norm,
        # scalars for the comparison table
        "mean_action_norm": float(per_step_norm.mean()) if per_step_norm.size else None,
        "max_action_norm": float(per_step_norm.max()) if per_step_norm.size else None,
        "mean_delta_norm": float(delta_norm.mean()) if delta_norm.size else 0.0,
        "gripper_min": float(gripper_all.min()) if gripper_all.size else None,
        "gripper_max": float(gripper_all.max()) if gripper_all.size else None,
    }


# --------------------------------------------------------------------------- #
# plotting
# --------------------------------------------------------------------------- #
def _save_fig(fig, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    fig.tight_layout()
    fig.savefig(path, dpi=110)
    plt.close(fig)


def make_plots(stats: dict, out_dir: Path, scenario_id: str) -> dict:
    """Write the 4 PNGs; return {logical_name: Path}. Requires matplotlib."""
    out_dir.mkdir(parents=True, exist_ok=True)
    calls = np.arange(stats["n_calls"])
    paths: dict[str, Path] = {}

    # 1) action norm over policy calls
    fig, ax = plt.subplots(figsize=(5.6, 3.2))
    ax.plot(calls, stats["step0_norm"], "o-", label="‖action‖ at chunk step 0")
    ax.plot(calls, stats["chunk_mean_norm"], "s--", label="‖mean action over 16-step chunk‖")
    ax.set_xlabel("policy call index")
    ax.set_ylabel("L2 norm")
    ax.set_title(f"{scenario_id}\naction L2 norm per policy call")
    ax.grid(alpha=0.3)
    ax.legend(fontsize=8)
    p = out_dir / "action_norm.png"
    _save_fig(fig, p)
    paths["action_norm"] = p

    # 2) mean action value per DoF (bar + std error bars)
    fig, ax = plt.subplots(figsize=(5.6, 3.2))
    x = np.arange(stats["n_dof"])
    ax.bar(x, stats["mean_per_dof"], yerr=stats["std_per_dof"], capsize=3, color="#4C78A8")
    ax.axhline(0.0, color="k", lw=0.6)
    ax.set_xticks(x)
    ax.set_xticklabels(stats["dof_names"], rotation=35, ha="right", fontsize=8)
    ax.set_ylabel("mean value (±std) over all calls×steps")
    gname = stats["dof_names"][stats["gripper_col"]]
    ax.set_title(f"{scenario_id}\nmean action per DoF  (gripper DoF = {gname}, col {stats['gripper_col']})")
    ax.grid(alpha=0.3, axis="y")
    p = out_dir / "action_mean_per_dof.png"
    _save_fig(fig, p)
    paths["action_mean_per_dof"] = p

    # 3) gripper command over time
    fig, ax = plt.subplots(figsize=(5.6, 3.2))
    ax.step(calls, stats["gripper_step0"], where="post", marker="o", color="#E45756")
    ax.set_ylim(min(-1.1, float(stats["gripper_step0"].min()) - 0.1) if stats["gripper_step0"].size else -1.1,
                max(1.1, float(stats["gripper_step0"].max()) + 0.1) if stats["gripper_step0"].size else 1.1)
    ax.set_xlabel("policy call index")
    ax.set_ylabel(f"{gname} (chunk step 0)")
    ax.set_title(f"{scenario_id}\ngripper command over time  (DoF '{gname}', col {stats['gripper_col']})")
    ax.grid(alpha=0.3)
    p = out_dir / "gripper_over_time.png"
    _save_fig(fig, p)
    paths["gripper_over_time"] = p

    # 4) action delta norm between consecutive chunks
    fig, ax = plt.subplots(figsize=(5.6, 3.2))
    if stats["delta_norm"].size:
        ax.plot(np.arange(1, stats["n_calls"]), stats["delta_norm"], "o-", color="#54A24B")
    else:
        ax.text(0.5, 0.5, "only one policy call\n(no consecutive delta)", ha="center", va="center",
                transform=ax.transAxes, fontsize=10)
    ax.set_xlabel("policy call index i  (delta between chunk i and i-1)")
    ax.set_ylabel("‖Δ chunk‖_F  (Frobenius over 16×n_dof)")
    ax.set_title(f"{scenario_id}\naction delta norm between consecutive chunks")
    ax.grid(alpha=0.3)
    p = out_dir / "action_delta_norm.png"
    _save_fig(fig, p)
    paths["action_delta_norm"] = p

    return paths


# --------------------------------------------------------------------------- #
# per-scenario processing
# --------------------------------------------------------------------------- #
def process_scenario(scenario_id: str, run_dir: Path, summary_row: dict | None,
                     make_plots_flag: bool, embed_video: bool) -> dict:
    """Return a dict with everything needed to render the card + table row."""
    sc_dir = run_dir / scenario_id
    warnings: list[str] = []
    rec: dict = {"scenario_id": scenario_id, "scenario_dir": sc_dir, "warnings": warnings}

    if not sc_dir.is_dir():
        warnings.append(f"scenario directory not found: {sc_dir}")
        rec.update(label=summary_row.get("label") if summary_row else None, seed=None,
                   instruction=None, success=summary_row.get("success") if summary_row else None,
                   max_episode_steps=None, num_policy_calls=None, actions_shape=None,
                   error=(summary_row or {}).get("error_if_any"), stats=None, plot_paths={},
                   video_path=None, frames_dir=None, n_frames=0,
                   files={}, video_data_uri=None)
        return rec

    meta, err = _load_json(sc_dir / "metadata.json")
    if err:
        warnings.append(err)
        meta = {}
    roll, err = _load_json(sc_dir / "rollout_summary.json")
    if err:
        warnings.append(err)
        roll = {}

    manifest_entry = (meta or {}).get("manifest_entry", {}) or {}
    resolved = (meta or {}).get("resolved", {}) or {}
    label = manifest_entry.get("label") or (meta or {}).get("label") or (summary_row or {}).get("label")
    seed = manifest_entry.get("seed", resolved.get("seed", (summary_row or {}).get("seed")))
    instruction = manifest_entry.get("instruction") or roll.get("instruction")
    env_name = resolved.get("env_name") or manifest_entry.get("env_name") or roll.get("env_name")
    max_episode_steps = (resolved.get("max_episode_steps")
                         or manifest_entry.get("max_episode_steps")
                         or (roll or {}).get("requested_max_episode_steps"))
    success = roll.get("success", (summary_row or {}).get("success"))
    error = (roll or {}).get("error") or (summary_row or {}).get("error_if_any")
    num_policy_calls = (roll or {}).get("n_get_action_calls") or (summary_row or {}).get("n_action_calls")
    obs_noise_std = resolved.get("obs_noise_std", manifest_entry.get("obs_noise_std"))
    action_repeat = resolved.get("action_repeat", manifest_entry.get("action_repeat"))
    notes = manifest_entry.get("notes")

    # actions.npy
    actions_shape = None
    stats = None
    plot_paths: dict[str, Path] = {}
    actions_path = sc_dir / "actions.npy"
    if actions_path.exists():
        try:
            arr = np.load(actions_path, allow_pickle=True)
            actions_shape = tuple(int(x) for x in np.asarray(arr).shape) if np.asarray(arr).dtype != object else f"object[{len(arr)}]"
            acts3 = _squeeze_actions(arr)
            stats = compute_action_stats(acts3)
            if num_policy_calls is None:
                num_policy_calls = stats["n_calls"]
            if make_plots_flag and _MPL_OK:
                try:
                    plot_paths = make_plots(stats, run_dir / "plots" / scenario_id, scenario_id)
                except Exception as e:  # noqa: BLE001
                    warnings.append(f"plot generation failed: {e!r}")
            elif make_plots_flag and not _MPL_OK:
                warnings.append(f"matplotlib unavailable, plots skipped: {_MPL_ERR}")
        except Exception as e:  # noqa: BLE001
            warnings.append(f"failed to load/parse actions.npy: {e!r}")
    else:
        warnings.append("missing file: actions.npy")

    # video
    video_path = sc_dir / "video.mp4"
    video_present = video_path.exists()
    if not video_present:
        warnings.append("missing file: video.mp4")
    video_data_uri = None
    if video_present and embed_video:
        video_data_uri = _b64_data_uri(video_path, "video/mp4")
        if video_data_uri is None:
            warnings.append("failed to base64-embed video.mp4")

    # frames/
    frames_dir = sc_dir / "frames"
    n_frames = 0
    first_frame_uri = None
    if frames_dir.is_dir():
        frame_files = sorted(frames_dir.glob("*.png"))
        n_frames = len(frame_files)
        if not video_present and frame_files:
            first_frame_uri = _b64_data_uri(frame_files[0], "image/png")
    # (frames/ is optional; only warn if neither video nor frames exist)
    if not video_present and n_frames == 0:
        warnings.append("no video.mp4 and no frames/ — nothing to preview")

    # raw-file links (present-or-not)
    files = {}
    for name in ("metadata.json", "rollout_summary.json", "stdout.log", "stderr.log"):
        p = sc_dir / name
        files[name] = {"rel": _rel(p, run_dir), "exists": p.exists()}
        if not p.exists():
            warnings.append(f"missing file: {name}")

    rec.update(
        label=label, seed=seed, instruction=instruction, env_name=env_name,
        success=success, max_episode_steps=max_episode_steps, num_policy_calls=num_policy_calls,
        actions_shape=actions_shape, error=error, obs_noise_std=obs_noise_std,
        action_repeat=action_repeat, notes=notes, stats=stats, plot_paths=plot_paths,
        video_path=video_path if video_present else None, video_present=video_present,
        video_data_uri=video_data_uri, frames_dir=frames_dir if frames_dir.is_dir() else None,
        n_frames=n_frames, first_frame_uri=first_frame_uri, files=files,
        output_dir=str((summary_row or {}).get("output_dir") or _rel(sc_dir, run_dir.parent.parent)),
        elapsed_sec=(roll or {}).get("elapsed_sec", (summary_row or {}).get("elapsed_sec")),
        video_saved=video_present,
    )
    return rec


# --------------------------------------------------------------------------- #
# HTML rendering
# --------------------------------------------------------------------------- #
_CSS = """
body{font-family:-apple-system,BlinkMacSystemFont,"Segoe UI",Roboto,Helvetica,Arial,sans-serif;
     margin:0;padding:24px;background:#f4f5f7;color:#1a1a1a;}
h1{margin:0 0 4px 0;} .sub{color:#666;margin-bottom:20px;font-size:14px;}
.meta-grid{display:grid;grid-template-columns:repeat(auto-fit,minmax(220px,1fr));gap:6px 18px;
           background:#fff;border:1px solid #e0e0e0;border-radius:8px;padding:14px 18px;margin-bottom:22px;font-size:13px;}
.meta-grid b{color:#444;}
table.cmp{border-collapse:collapse;width:100%;background:#fff;font-size:12.5px;margin-bottom:28px;
          box-shadow:0 1px 3px rgba(0,0,0,.08);border-radius:8px;overflow:hidden;}
table.cmp th,table.cmp td{border-bottom:1px solid #eee;padding:7px 10px;text-align:left;white-space:nowrap;}
table.cmp th{background:#2d3748;color:#fff;position:sticky;top:0;}
table.cmp tr:hover{background:#f7f9fc;}
.badge{display:inline-block;padding:1px 8px;border-radius:10px;font-size:11px;font-weight:600;}
.b-normal{background:#e3f0ff;color:#1a5fb4;} .b-abn{background:#ffe9d6;color:#b35a00;}
.b-ok{background:#e6f6ea;color:#1a7f37;} .b-fail{background:#fde8e8;color:#b42318;} .b-unk{background:#eee;color:#555;}
.card{background:#fff;border:1px solid #e0e0e0;border-radius:10px;padding:18px;margin-bottom:22px;
      box-shadow:0 1px 3px rgba(0,0,0,.06);}
.card h2{margin:0 0 2px 0;font-size:18px;} .card .cmeta{color:#555;font-size:13px;margin-bottom:10px;}
.card .cols{display:grid;grid-template-columns:minmax(280px,360px) 1fr;gap:18px;}
.kv{font-size:13px;line-height:1.7;} .kv b{color:#444;display:inline-block;min-width:130px;}
.plots{display:grid;grid-template-columns:repeat(auto-fit,minmax(300px,1fr));gap:10px;}
.plots img{width:100%;border:1px solid #eee;border-radius:6px;background:#fff;}
video{width:100%;max-width:340px;border-radius:6px;background:#000;}
.warn{background:#fff8e1;border:1px solid #ffe082;border-radius:6px;padding:8px 12px;margin:8px 0;font-size:12.5px;color:#7a5c00;}
.err{background:#fdecea;border:1px solid #f5c6c0;border-radius:6px;padding:8px 12px;margin:8px 0;font-size:12.5px;color:#7a1f17;font-family:monospace;white-space:pre-wrap;}
.files a{margin-right:12px;font-size:12.5px;} .files .missing{color:#aaa;text-decoration:line-through;}
.notes{font-size:12.5px;color:#555;font-style:italic;margin-top:8px;}
code{background:#f0f0f0;padding:1px 4px;border-radius:3px;font-size:12px;}
"""


def _badge_label(label: str | None) -> str:
    if label == "normal":
        return '<span class="badge b-normal">normal</span>'
    if label == "abnormal_probe":
        return '<span class="badge b-abn">abnormal_probe</span>'
    return f'<span class="badge b-unk">{html.escape(str(label))}</span>'


def _badge_success(s) -> str:
    if s is True:
        return '<span class="badge b-ok">success: true</span>'
    if s is False:
        return '<span class="badge b-fail">success: false</span>'
    return '<span class="badge b-unk">success: unknown</span>'


def render_html(run_dir: Path, summary: dict | None, records: list[dict], table_rows: list[dict],
                args) -> str:
    parts: list[str] = []
    parts.append("<!DOCTYPE html><html><head><meta charset='utf-8'>")
    parts.append(f"<title>LIBERO smoke-test visual report — {html.escape(run_dir.name)}</title>")
    parts.append(f"<style>{_CSS}</style></head><body>")
    parts.append("<h1>LIBERO smoke-test visual review</h1>")
    parts.append(f"<div class='sub'>run dir: <code>{html.escape(str(run_dir))}</code> &nbsp;|&nbsp; "
                 f"generated {html.escape(_dt.datetime.now().isoformat(timespec='seconds'))} &nbsp;|&nbsp; "
                 f"scenarios: {len(records)}</div>")

    # run-level metadata
    if summary:
        t = summary.get("totals", {})
        parts.append("<div class='meta-grid'>")
        for k, v in [
            ("model_path", summary.get("model_path")),
            ("policy_server", f"{summary.get('policy_server',{}).get('host')}:{summary.get('policy_server',{}).get('port')}"),
            ("manifest", summary.get("manifest")),
            ("server_command", summary.get("server_command")),
            ("n_scenarios", summary.get("n_scenarios")),
            ("n_normal / n_abnormal_probe", f"{summary.get('n_normal')} / {summary.get('n_abnormal_probe')}"),
            ("rollout_started / actions_produced / video_saved / errors",
             f"{t.get('rollout_started')} / {t.get('actions_produced')} / {t.get('video_saved')} / {t.get('errors')}"),
            ("run timestamp", summary.get("timestamp")),
            ("dry_run", summary.get("dry_run")),
        ]:
            parts.append(f"<div><b>{html.escape(str(k))}:</b> {html.escape(_fmt(v))}</div>")
        parts.append("</div>")
    if not _MPL_OK:
        parts.append(f"<div class='warn'>matplotlib not available ({html.escape(str(_MPL_ERR))}) — plots were skipped.</div>")
    if args.no_plots:
        parts.append("<div class='warn'>--no-plots was set — no PNG plots were generated this run.</div>")

    # comparison table
    parts.append("<h2>Comparison across all scenarios</h2>")
    cols = ["scenario_id", "label", "seed", "success", "num_policy_calls", "video_saved",
            "mean_action_norm", "max_action_norm", "mean_delta_norm", "gripper_min", "gripper_max", "output_dir"]
    parts.append("<table class='cmp'><thead><tr>" + "".join(f"<th>{html.escape(c)}</th>" for c in cols) + "</tr></thead><tbody>")
    for row in table_rows:
        parts.append("<tr>")
        for c in cols:
            v = row.get(c)
            if c == "label":
                cell = _badge_label(v)
            elif c == "success":
                cell = _badge_success(v)
            elif c == "scenario_id":
                cell = f"<a href='#{html.escape(str(v))}'>{html.escape(str(v))}</a>"
            elif c in ("mean_action_norm", "max_action_norm", "mean_delta_norm", "gripper_min", "gripper_max"):
                cell = "—" if v is None else f"{float(v):.4g}"
            else:
                cell = html.escape(_fmt(v))
            parts.append(f"<td>{cell}</td>")
        parts.append("</tr>")
    parts.append("</tbody></table>")

    # one card per scenario
    for rec in records:
        sid = rec["scenario_id"]
        parts.append(f"<div class='card' id='{html.escape(sid)}'>")
        parts.append(f"<h2>{html.escape(sid)}</h2>")
        parts.append("<div class='cmeta'>" + _badge_label(rec.get("label")) + " " + _badge_success(rec.get("success"))
                     + f" &nbsp; seed={html.escape(_fmt(rec.get('seed')))}"
                     + (f" &nbsp; env=<code>{html.escape(str(rec.get('env_name')))}</code>" if rec.get("env_name") else "")
                     + "</div>")
        if rec.get("error"):
            parts.append(f"<div class='err'>error: {html.escape(str(rec['error']))}</div>")
        for w in rec.get("warnings", []):
            parts.append(f"<div class='warn'>⚠ {html.escape(str(w))}</div>")

        parts.append("<div class='cols'>")

        # left column: key/values + video + file links
        parts.append("<div>")
        st = rec.get("stats") or {}
        kv = [
            ("scenario_id", sid),
            ("label", rec.get("label")),
            ("seed", rec.get("seed")),
            ("instruction", rec.get("instruction")),
            ("success", rec.get("success")),
            ("max_episode_steps", rec.get("max_episode_steps")),
            ("num policy calls", rec.get("num_policy_calls")),
            ("actions.npy shape", rec.get("actions_shape")),
            ("action horizon × DoF", f"{st.get('horizon')} × {st.get('n_dof')}" if st else "—"),
            ("gripper DoF (col)", (f"{st['dof_names'][st['gripper_col']]} (col {st['gripper_col']})" if st else "—")),
            ("mean / max ‖action‖", f"{_fmt(st.get('mean_action_norm'))} / {_fmt(st.get('max_action_norm'))}" if st else "—"),
            ("mean ‖Δ chunk‖", _fmt(st.get("mean_delta_norm")) if st else "—"),
            ("gripper min / max", f"{_fmt(st.get('gripper_min'))} / {_fmt(st.get('gripper_max'))}" if st else "—"),
            ("obs_noise_std", rec.get("obs_noise_std")),
            ("action_repeat", rec.get("action_repeat")),
            ("elapsed_sec", rec.get("elapsed_sec")),
            ("frames", (f"{rec.get('n_frames')} PNGs" if rec.get("n_frames") else "none")),
        ]
        parts.append("<div class='kv'>")
        for k, v in kv:
            parts.append(f"<div><b>{html.escape(str(k))}:</b> {html.escape(_fmt(v))}</div>")
        parts.append("</div>")

        # video preview
        parts.append("<div style='margin-top:10px'>")
        if rec.get("video_data_uri"):
            parts.append(f"<video controls preload='metadata' src='{rec['video_data_uri']}'></video>")
        elif rec.get("video_present"):
            vrel = _rel(rec["video_path"], run_dir)
            parts.append(f"<video controls preload='metadata' src='{html.escape(vrel)}'></video>"
                         f"<div style='font-size:12px'><a href='{html.escape(vrel)}'>{html.escape(vrel)}</a></div>")
        elif rec.get("first_frame_uri"):
            parts.append(f"<img src='{rec['first_frame_uri']}' style='max-width:340px;border-radius:6px'/>"
                         f"<div class='warn'>no video.mp4 — showing first frame only</div>")
        else:
            parts.append("<div class='warn'>no video / frames available to preview</div>")
        parts.append("</div>")

        # file links
        parts.append("<div class='files' style='margin-top:10px'>")
        for name, info in (rec.get("files") or {}).items():
            if info["exists"]:
                parts.append(f"<a href='{html.escape(info['rel'])}'>{html.escape(name)}</a>")
            else:
                parts.append(f"<span class='missing'>{html.escape(name)}</span>")
        if rec.get("frames_dir"):
            parts.append(f"<a href='{html.escape(_rel(rec['frames_dir'], run_dir))}/'>frames/ ({rec.get('n_frames')})</a>")
        parts.append("</div>")
        if rec.get("notes"):
            parts.append(f"<div class='notes'>note: {html.escape(str(rec['notes']))}</div>")
        parts.append("</div>")  # end left column

        # right column: plots
        parts.append("<div>")
        if rec.get("plot_paths"):
            parts.append("<div class='plots'>")
            for logical in ("action_norm", "action_mean_per_dof", "gripper_over_time", "action_delta_norm"):
                p = rec["plot_paths"].get(logical)
                if p and Path(p).exists():
                    parts.append(f"<a href='{html.escape(_rel(Path(p), run_dir))}'>"
                                 f"<img src='{html.escape(_rel(Path(p), run_dir))}' alt='{logical}'/></a>")
                else:
                    parts.append(f"<div class='warn'>plot '{logical}.png' not available</div>")
            parts.append("</div>")
        elif rec.get("stats") is None:
            parts.append("<div class='warn'>no actions.npy — no action statistics / plots</div>")
        else:
            parts.append("<div class='warn'>plots not generated (use without --no-plots, and ensure matplotlib is installed)</div>")
        parts.append("</div>")  # end right column

        parts.append("</div>")  # end cols
        parts.append("</div>")  # end card

    parts.append("</body></html>")
    return "\n".join(parts)


# --------------------------------------------------------------------------- #
# main
# --------------------------------------------------------------------------- #
def _discover_scenarios(run_dir: Path, summary: dict | None) -> tuple[list[str], dict[str, dict]]:
    rows_by_id: dict[str, dict] = {}
    order: list[str] = []
    if summary and isinstance(summary.get("scenarios"), list):
        for r in summary["scenarios"]:
            sid = r.get("scenario_id")
            if sid:
                rows_by_id[sid] = r
                order.append(sid)
    # add any scenario subdirs not already listed
    skip = {"plots"}
    for p in sorted(run_dir.iterdir()):
        if p.is_dir() and p.name not in skip and p.name not in order:
            order.append(p.name)
    return order, rows_by_id


def write_table_outputs(run_dir: Path, table_rows: list[dict]) -> tuple[Path, Path]:
    cols = ["scenario_id", "label", "seed", "success", "num_policy_calls", "video_saved",
            "mean_action_norm", "max_action_norm", "mean_delta_norm", "gripper_min", "gripper_max", "output_dir"]
    csv_path = run_dir / "visual_summary.csv"
    with open(csv_path, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=cols)
        w.writeheader()
        for row in table_rows:
            w.writerow({c: ("" if row.get(c) is None else row.get(c)) for c in cols})
    json_path = run_dir / "visual_summary.json"
    with open(json_path, "w") as f:
        json.dump({"run_dir": str(run_dir), "n_scenarios": len(table_rows),
                   "generated": _dt.datetime.now().isoformat(timespec="seconds"),
                   "columns": cols, "rows": table_rows}, f, indent=2)
    return csv_path, json_path


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--run-dir", required=True, help="Path to the completed run dir (contains summary.json + scenario subdirs).")
    ap.add_argument("--open", action="store_true", help="Open the generated HTML report in a browser.")
    ap.add_argument("--no-plots", action="store_true", help="Do not (re)generate the PNG plots.")
    ap.add_argument("--no-video-embed", action="store_true", help="Link videos by relative path instead of base64-embedding them in the HTML.")
    args = ap.parse_args()

    run_dir = Path(args.run_dir).resolve()
    if not run_dir.is_dir():
        print(f"ERROR: --run-dir does not exist or is not a directory: {run_dir}", file=sys.stderr)
        return 2

    summary, serr = _load_json(run_dir / "summary.json")
    if serr:
        print(f"WARNING: {serr} — proceeding by scanning scenario subdirectories.", file=sys.stderr)
    # summary.md is read only to confirm it exists / surface it; we don't parse it.
    summary_md_present = (run_dir / "summary.md").exists()
    if not summary_md_present:
        print("WARNING: summary.md not found (continuing).", file=sys.stderr)

    order, rows_by_id = _discover_scenarios(run_dir, summary)
    if not order:
        print(f"ERROR: no scenarios found under {run_dir}", file=sys.stderr)
        return 2
    print(f"Found {len(order)} scenario(s) in {run_dir}")

    make_plots_flag = not args.no_plots
    embed_video = not args.no_video_embed

    records: list[dict] = []
    table_rows: list[dict] = []
    all_warnings: list[str] = []
    n_videos_embedded = 0
    n_videos_present = 0
    for sid in order:
        try:
            rec = process_scenario(sid, run_dir, rows_by_id.get(sid), make_plots_flag, embed_video)
        except Exception as e:  # noqa: BLE001
            tb = traceback.format_exc()
            print(f"WARNING: scenario {sid} raised {e!r}; recording as failed card.", file=sys.stderr)
            rec = {"scenario_id": sid, "warnings": [f"internal error: {e!r}", tb], "stats": None,
                   "plot_paths": {}, "files": {}, "label": None, "seed": None, "success": None,
                   "video_present": False, "video_data_uri": None, "n_frames": 0}
        records.append(rec)
        all_warnings += [f"[{sid}] {w}" for w in rec.get("warnings", [])]
        if rec.get("video_present"):
            n_videos_present += 1
        if rec.get("video_data_uri"):
            n_videos_embedded += 1
        st = rec.get("stats") or {}
        table_rows.append({
            "scenario_id": sid,
            "label": rec.get("label"),
            "seed": rec.get("seed"),
            "success": rec.get("success"),
            "num_policy_calls": rec.get("num_policy_calls"),
            "video_saved": bool(rec.get("video_present")),
            "mean_action_norm": st.get("mean_action_norm"),
            "max_action_norm": st.get("max_action_norm"),
            "mean_delta_norm": st.get("mean_delta_norm"),
            "gripper_min": st.get("gripper_min"),
            "gripper_max": st.get("gripper_max"),
            "output_dir": rec.get("output_dir") or str(run_dir / sid),
        })

    csv_path, json_path = write_table_outputs(run_dir, table_rows)
    html_str = render_html(run_dir, summary, records, table_rows, args)
    html_path = run_dir / "visual_report.html"
    html_path.write_text(html_str, encoding="utf-8")

    # console summary
    print(f"\nWrote: {html_path}")
    print(f"Wrote: {csv_path}")
    print(f"Wrote: {json_path}")
    if make_plots_flag and _MPL_OK:
        print(f"Plots: {run_dir / 'plots'}/<scenario_id>/*.png")
    print(f"Videos: {n_videos_present}/{len(order)} present; "
          f"{'all ' if n_videos_embedded == len(order) and embed_video else ''}"
          f"{n_videos_embedded}/{len(order)} embedded in HTML"
          + ("" if embed_video else "  (--no-video-embed: linked by path instead)"))
    if all_warnings:
        print(f"\n{len(all_warnings)} warning(s):")
        for w in all_warnings:
            print(f"  - {w}")
    else:
        print("\nNo missing files or warnings.")

    if args.open:
        import webbrowser

        url = html_path.as_uri()
        opened = False
        try:
            opened = webbrowser.open(url)
        except Exception:
            opened = False
        if opened:
            print(f"\nOpened {url} in a browser.")
        else:
            print(f"\nCould not auto-open a browser. Open this file manually:\n  {html_path}")

    return 0


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