app.py

"""
Side-by-side viewer for forward-dynamics rollouts vs. ground truth.

Point it at any HF dataset repo laid out as:

    windows_manifest.json               # optional: list of window dicts
    metrics_summary.json                # optional: per-window MSE / PSNR
    windows/<episode_id>/gt.mp4         # ground-truth window
    windows/<episode_id>/generated.mp4  # model rollout for the same slice

Ground truth and generated are composed into a single video so the two panels
stay frame-locked; two independent players drift apart and cannot be scrubbed
together. Videos may stack N cameras vertically; N is inferred from the frame
aspect ratio and the camera selector adapts.

The app starts empty and assumes no dataset: type a repo id and press Load.
Nothing is downloaded at build time and no repo is ever pulled whole. Each
browser session gets its own scratch directory, keeps at most
MAX_WINDOWS_ON_DISK windows, and is deleted when the session ends.

Private repos need an HF_TOKEN secret with read access.
"""

from __future__ import annotations

import json
import os
import re
import shutil
import tempfile
import threading
import time
import uuid
from collections import OrderedDict
from pathlib import Path

import cv2
import gradio as gr
import imageio_ffmpeg
import numpy as np
from huggingface_hub import HfApi, RepoFolder, hf_hub_download

HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") or None

BASE_DIR = Path(tempfile.gettempdir()) / "wmviz"
BASE_DIR.mkdir(parents=True, exist_ok=True)

MAX_WINDOWS_ON_DISK = int(os.environ.get("MAX_WINDOWS_ON_DISK", 12))
MAX_PAIRS_IN_RAM = 2
MAX_COMPOSED_IN_RAM = 1
MAX_DISCOVER = int(os.environ.get("MAX_DISCOVER", 5000))
SESSION_TTL = int(os.environ.get("SESSION_TTL", 3600))

THREE_CAM_NAMES = ["Top", "Left wrist", "Right wrist"]
ALL_CAMS = -1
NO_CRIT = -1

FONT = cv2.FONT_HERSHEY_SIMPLEX
HDR, FTR, SEP = 26, 26, 4
SCALE = 2  # source frames are small; upscale so overlays stay legible

GT_COLOR = (130, 225, 140)
GEN_COLOR = (120, 180, 255)
DIFF_COLOR = (215, 215, 215)
CRIT_COLOR = (255, 95, 95)

SORTS = [
    "Hardest first (critical PSNR up)",
    "Easiest first (critical PSNR down)",
    "Task A-Z",
    "Manifest order",
]


# --------------------------------------------------------------------------- session store


def parse_repo(text: str) -> tuple[str, str | None]:
    """Split "owner/name@revision" into (repo_id, revision). Tolerates a pasted URL."""
    text = (text or "").strip().rstrip("/")
    m = re.match(r"^https?://huggingface\.co/(?:datasets/)?(.+?)(?:/tree/([^/]+))?$", text)
    if m:
        return m.group(1), m.group(2)
    if "@" in text:
        repo, _, rev = text.partition("@")
        return repo.strip(), rev.strip() or None
    return text, None


class Session:
    """One browser session's scratch space: bounded on disk and in RAM."""

    def __init__(self, sid: str):
        self.sid = sid
        self.root = BASE_DIR / sid
        self.data = self.root / "data"
        self.render = self.root / "render"
        for d in (self.data, self.render):
            d.mkdir(parents=True, exist_ok=True)
        self.repo = ""
        self.revision: str | None = None
        self.rows: list[dict] = []
        self._on_disk: OrderedDict[str, bool] = OrderedDict()
        self._pairs: OrderedDict[str, tuple[np.ndarray, np.ndarray]] = OrderedDict()
        self._composed: OrderedDict[tuple, np.ndarray] = OrderedDict()
        self.touched = time.time()

    def fetch(self, rel: str, optional: bool = False) -> str | None:
        """Download one repo file into this session's directory."""
        self.touched = time.time()
        try:
            return hf_hub_download(
                repo_id=self.repo,
                filename=rel,
                repo_type="dataset",
                revision=self.revision,
                token=HF_TOKEN,
                local_dir=str(self.data),
            )
        except Exception:
            if optional:
                return None
            raise

    def window_file(self, episode_id: str, name: str, optional: bool = False) -> str | None:
        path = self.fetch(f"windows/{episode_id}/{name}", optional=optional)
        if path:
            self._keep(episode_id)
        return path

    def _keep(self, episode_id: str) -> None:
        """Register a window as resident and evict the least recently used ones."""
        self._on_disk[episode_id] = True
        self._on_disk.move_to_end(episode_id)
        while len(self._on_disk) > MAX_WINDOWS_ON_DISK:
            old, _ = self._on_disk.popitem(last=False)
            shutil.rmtree(self.data / "windows" / old, ignore_errors=True)
            self._pairs.pop(old, None)
            for key in [k for k in self._composed if k[0] == old]:
                self._composed.pop(key, None)
            for f in self.render.glob(f"{old}_*"):
                f.unlink(missing_ok=True)

    def disk_windows(self) -> int:
        return len(self._on_disk)

    def pair(self, episode_id: str) -> tuple[np.ndarray, np.ndarray]:
        hit = self._pairs.get(episode_id)
        if hit is not None:
            self._pairs.move_to_end(episode_id)
            return hit
        gt = read_video(self.window_file(episode_id, "gt.mp4"))
        gen = read_video(self.window_file(episode_id, "generated.mp4"))
        t = min(len(gt), len(gen))
        value = (gt[:t], gen[:t])
        self._pairs[episode_id] = value
        while len(self._pairs) > MAX_PAIRS_IN_RAM:
            self._pairs.popitem(last=False)
        return value

    def composed(self, episode_id: str, view: int, show_diff: bool, crit: int) -> np.ndarray:
        key = (episode_id, view, show_diff, crit)
        hit = self._composed.get(key)
        if hit is not None:
            self._composed.move_to_end(key)
            return hit
        value = compose(self.pair(episode_id), view, show_diff, crit)
        self._composed[key] = value
        while len(self._composed) > MAX_COMPOSED_IN_RAM:
            self._composed.popitem(last=False)
        return value

    def set_repo(self, repo: str, revision: str | None) -> None:
        if (repo, revision) != (self.repo, self.revision):
            self.reset()
            self.repo, self.revision = repo, revision

    def reset(self) -> None:
        self._on_disk.clear()
        self._pairs.clear()
        self._composed.clear()
        self.rows = []
        for d in (self.data, self.render):
            shutil.rmtree(d, ignore_errors=True)
            d.mkdir(parents=True, exist_ok=True)

    def close(self) -> None:
        self._pairs.clear()
        self._composed.clear()
        shutil.rmtree(self.root, ignore_errors=True)


SESSIONS: dict[str, Session] = {}
_LOCK = threading.Lock()


def sweep() -> None:
    """Drop sessions whose browser tab went away without a clean callback."""
    now = time.time()
    with _LOCK:
        stale = [s for s in SESSIONS.values() if now - s.touched > SESSION_TTL]
        for s in stale:
            SESSIONS.pop(s.sid, None)
            s.close()
        live = set(SESSIONS)
    for d in BASE_DIR.iterdir():  # dirs left behind by a previous process
        try:
            if d.is_dir() and d.name not in live and now - d.stat().st_mtime > SESSION_TTL:
                shutil.rmtree(d, ignore_errors=True)
        except OSError:
            pass


def get_session(sid: str | None) -> Session:
    with _LOCK:
        s = SESSIONS.get(sid or "")
        if s is None:
            sid = sid or uuid.uuid4().hex
            s = SESSIONS[sid] = Session(sid)
    s.touched = time.time()
    return s


def release(sid: str | None) -> None:
    """gr.State delete_callback - fires when the session ends or its TTL expires."""
    with _LOCK:
        s = SESSIONS.pop(sid or "", None)
    if s is not None:
        s.close()


# --------------------------------------------------------------------------- index


def _row(order: int, episode_id: str, w: dict, m: dict) -> dict:
    return {
        "order": order,
        "episode_id": episode_id,
        "task": w.get("task") or "?",
        "edge": w.get("edge_type") or "?",
        "arm": w.get("arm") or "?",
        "critical_frame": w.get("critical_frame"),
        "start_frame": w.get("start_frame"),
        "full_psnr": m.get("full_psnr"),
        "crit_psnr": m.get("crit_psnr"),
        "full_mse": m.get("full_mse"),
        "crit_mse": m.get("crit_mse"),
    }


def discover_episodes(session: Session) -> tuple[list[str], bool]:
    """List windows/ one level deep, capped. Returns (ids, was_truncated)."""
    tree = HfApi().list_repo_tree(
        session.repo,
        path_in_repo="windows",
        recursive=False,
        revision=session.revision,
        repo_type="dataset",
        token=HF_TOKEN,
    )
    eps, truncated = [], False
    for entry in tree:
        if not isinstance(entry, RepoFolder):
            continue
        eps.append(entry.path.split("/")[-1])
        if len(eps) >= MAX_DISCOVER:
            truncated = True
            break
    return sorted(eps), truncated


def build_index(session: Session) -> str:
    """Populate session.rows from the manifest, or by listing the repo. Returns a note."""
    metrics = []
    path = session.fetch("metrics_summary.json", optional=True)
    if path:
        try:
            metrics = json.load(open(path))
        except Exception:
            metrics = []
    by_id = {m.get("id"): m for m in metrics if isinstance(m, dict)}

    manifest = None
    path = session.fetch("windows_manifest.json", optional=True)
    if path:
        try:
            loaded = json.load(open(path))
            if isinstance(loaded, list) and loaded:
                manifest = loaded
        except Exception:
            manifest = None

    if manifest is not None:
        session.rows = [
            _row(i, w.get("episode_id", f"window_{i}"), w, by_id.get(w.get("episode_id"), {}))
            for i, w in enumerate(manifest)
        ]
        note = "`windows_manifest.json`"
    else:
        eps, truncated = discover_episodes(session)
        if not eps:
            raise FileNotFoundError(
                "no windows/<episode_id>/ directories found - is this the right repo?"
            )
        session.rows = [_row(i, ep, {}, by_id.get(ep, {})) for i, ep in enumerate(eps)]
        note = "directory listing (no `windows_manifest.json`)"
        if truncated:
            note += f", capped at {MAX_DISCOVER}"

    if metrics:
        note += " + `metrics_summary.json`"
    return note


def short_id(episode_id: str) -> str:
    return re.sub(r"^(episode|window|ep)[-_]", "", episode_id)[:8]


def label_for(row: dict) -> str:
    psnr = row["crit_psnr"]
    head = f"{psnr:5.1f} dB" if isinstance(psnr, (int, float)) else "  -- dB"
    tags = " ".join(t for t in (row["edge"], row["arm"]) if t and t != "?")
    parts = [head] + ([tags] if tags else []) + [row["task"], short_id(row["episode_id"])]
    return " | ".join(parts)


def filter_rows(session: Session, query: str, sort: str) -> list[dict]:
    rows = session.rows
    q = (query or "").strip().lower()
    if q:
        rows = [r for r in rows if q in r["task"].lower() or q in r["episode_id"].lower()]

    def psnr(r, missing):
        v = r["crit_psnr"]
        return v if isinstance(v, (int, float)) else missing

    if sort == SORTS[0]:
        return sorted(rows, key=lambda r: psnr(r, float("inf")))
    if sort == SORTS[1]:
        return sorted(rows, key=lambda r: -psnr(r, float("-inf")))
    if sort == SORTS[2]:
        return sorted(rows, key=lambda r: (r["task"], r["episode_id"]))
    return sorted(rows, key=lambda r: r["order"])


def row_by_id(session: Session, episode_id: str) -> dict:
    for r in session.rows:
        if r["episode_id"] == episode_id:
            return r
    return _row(0, episode_id, {}, {})


# --------------------------------------------------------------------------- video


def read_video(path: str) -> np.ndarray:
    """Decode an mp4 into a uint8 [T, H, W, 3] RGB array."""
    cap = cv2.VideoCapture(path)
    frames = []
    while True:
        ok, frame = cap.read()
        if not ok:
            break
        frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
    cap.release()
    if not frames:
        raise RuntimeError(f"no frames decoded from {path}")
    return np.stack(frames)


def n_cameras(frames: np.ndarray) -> int:
    """How many roughly-square cameras are stacked vertically in the frame."""
    h, w = frames.shape[1], frames.shape[2]
    return max(1, min(6, int(round(h / w))))


def camera_names(n: int) -> list[str]:
    return THREE_CAM_NAMES if n == 3 else [f"Camera {i + 1}" for i in range(n)]


def camera_choices(n: int) -> list[tuple[str, int]]:
    choices = [(name, i) for i, name in enumerate(camera_names(n))]
    return ([("All cameras", ALL_CAMS)] + choices) if n > 1 else choices


def cam_slices(frames: np.ndarray, view: int) -> list[np.ndarray]:
    n = n_cameras(frames)
    h = frames.shape[1] // n
    rows = range(n) if view == ALL_CAMS or not 0 <= view < n else [view]
    return [frames[:, i * h : (i + 1) * h] for i in rows]


def letterbox_rows(*bands: np.ndarray, thresh: int = 16) -> tuple[int, int]:
    """Rows to keep after dropping the black bars shared by every given band."""
    h = bands[0].shape[1]
    lit = np.zeros(h, bool)
    for b in bands:
        lit |= b.reshape(b.shape[0], h, -1).max(axis=(0, 2)) > thresh
    idx = np.flatnonzero(lit)
    if idx.size < h * 0.1:  # nothing meaningful to trim
        return 0, h
    return int(idx[0]), int(idx[-1]) + 1


def crop_view(gt: np.ndarray, gen: np.ndarray, view: int) -> tuple[np.ndarray, np.ndarray]:
    """Select cameras and trim the letterbox, identically for both videos."""
    gt_out, gen_out = [], []
    for g, p in zip(cam_slices(gt, view), cam_slices(gen, view)):
        lo, hi = letterbox_rows(g, p)
        gt_out.append(g[:, lo:hi])
        gen_out.append(p[:, lo:hi])
    return np.concatenate(gt_out, axis=1), np.concatenate(gen_out, axis=1)


def upscale(frames: np.ndarray, scale: int) -> np.ndarray:
    if scale == 1:
        return frames
    return np.stack(
        [cv2.resize(f, None, fx=scale, fy=scale, interpolation=cv2.INTER_NEAREST) for f in frames]
    )


def panel(frames: np.ndarray, title: str, color: tuple[int, int, int]) -> np.ndarray:
    t, h, w, _ = frames.shape
    out = np.zeros((t, HDR + h, w, 3), np.uint8)
    out[:, HDR:] = frames
    bar = np.full((HDR, w, 3), 22, np.uint8)
    cv2.putText(bar, title, (8, HDR - 8), FONT, 0.52, color, 1, cv2.LINE_AA)
    out[:, :HDR] = bar
    return out


def diff_frames(gt: np.ndarray, gen: np.ndarray, gain: float = 3.0) -> np.ndarray:
    d = np.abs(gt.astype(np.int16) - gen.astype(np.int16)).mean(axis=-1)
    d = np.clip(d * gain, 0, 255).astype(np.uint8)
    return np.stack([cv2.applyColorMap(f, cv2.COLORMAP_INFERNO)[:, :, ::-1] for f in d])


def compose(pair: tuple[np.ndarray, np.ndarray], view: int, show_diff: bool, crit_idx: int) -> np.ndarray:
    """Build the labelled side-by-side strip: [T, H, W, 3] RGB."""
    gt_c, gen_c = crop_view(pair[0], pair[1], view)
    gt_v, gen_v = upscale(gt_c, SCALE), upscale(gen_c, SCALE)

    panels = [panel(gt_v, "GROUND TRUTH", GT_COLOR), panel(gen_v, "GENERATED", GEN_COLOR)]
    if show_diff:
        panels.append(panel(diff_frames(gt_v, gen_v), "|DIFF| x3", DIFF_COLOR))

    t, ph, _, _ = panels[0].shape
    gap = np.zeros((t, ph, SEP, 3), np.uint8)
    stitched = [panels[0]]
    for p in panels[1:]:
        stitched += [gap, p]
    comp = np.concatenate(stitched, axis=2)

    t, h, w, _ = comp.shape
    out = np.zeros((t, h + FTR, w, 3), np.uint8)
    out[:, :h] = comp
    has_crit = 0 <= crit_idx < t
    tl0, tl1, y = int(w * 0.45), w - 14, FTR // 2
    for i in range(t):
        bar = np.full((FTR, w, 3), 22, np.uint8)
        is_crit = has_crit and i == crit_idx
        text = f"frame {i:02d}/{t - 1}" + ("   CRITICAL FRAME" if is_crit else "")
        cv2.putText(bar, text, (8, FTR - 8), FONT, 0.5, CRIT_COLOR if is_crit else (200, 200, 200), 1, cv2.LINE_AA)
        cv2.line(bar, (tl0, y), (tl1, y), (70, 70, 70), 2, cv2.LINE_AA)
        if has_crit:
            cx = int(tl0 + (tl1 - tl0) * crit_idx / max(t - 1, 1))
            cv2.line(bar, (cx, y - 7), (cx, y + 7), CRIT_COLOR, 2, cv2.LINE_AA)
        px = int(tl0 + (tl1 - tl0) * i / max(t - 1, 1))
        cv2.circle(bar, (px, y), 4, (255, 255, 255), -1, cv2.LINE_AA)
        out[i, h:] = bar
        if is_crit:
            cv2.rectangle(out[i], (0, 0), (w - 1, h + FTR - 1), CRIT_COLOR, 3)

    # yuv420p needs even dimensions
    ph, pw = out.shape[1] % 2, out.shape[2] % 2
    if ph or pw:
        out = np.pad(out, ((0, 0), (0, ph), (0, pw), (0, 0)))
    return out


def write_mp4(frames: np.ndarray, path: Path, fps: float) -> str:
    _, h, w, _ = frames.shape
    writer = imageio_ffmpeg.write_frames(
        str(path),
        size=(w, h),
        fps=fps,
        codec="libx264",
        pix_fmt_in="rgb24",
        pix_fmt_out="yuv420p",
        macro_block_size=1,
        output_params=["-crf", "16", "-preset", "veryfast", "-movflags", "+faststart"],
    )
    writer.send(None)
    for f in frames:
        writer.send(np.ascontiguousarray(f).tobytes())
    writer.close()
    return str(path)


def render_video(session: Session, episode_id: str, view: int, show_diff: bool, fps: float, crit: int) -> str:
    path = session.render / f"{episode_id}_{view}_{int(show_diff)}_{fps:g}.mp4"
    if not path.exists():
        write_mp4(session.composed(episode_id, view, show_diff, crit), path, fps)
    return str(path)


# --------------------------------------------------------------------------- ui glue


def metrics_md(row: dict) -> str:
    def fmt(v, unit=""):
        return f"{v:.2f}{unit}" if isinstance(v, (int, float)) else "-"

    out = [
        f"**{row['task'].replace('_', ' ')}**",
        "",
        f"`{row['episode_id']}`",
        "",
        "| | full window | critical frame |",
        "|---|---|---|",
        f"| **PSNR** | {fmt(row['full_psnr'], ' dB')} | {fmt(row['crit_psnr'], ' dB')} |",
        f"| **MSE** | {fmt(row['full_mse'])} | {fmt(row['crit_mse'])} |",
        "",
    ]
    bits = [b for b in (row["edge"], row["arm"]) if b and b != "?"]
    if row["critical_frame"] is not None:
        bits.append(f"critical frame {row['critical_frame']}")
    if bits:
        out.append(" | ".join(bits))
    return "\n".join(out)


def critical_index(row: dict) -> int:
    start, crit = row.get("start_frame"), row.get("critical_frame")
    if start is None or crit is None:
        return NO_CRIT
    return int(crit) - int(start)


def on_load(repo_text, sid):
    """Resolve a repo and load only its index - no window data is pulled here."""
    sweep()
    session = get_session(sid)
    repo, revision = parse_repo(repo_text)
    if not repo:
        return session.sid, "Enter a dataset repo id, e.g. `owner/name`."
    session.set_repo(repo, revision)
    try:
        note = build_index(session)
    except Exception as exc:
        session.rows = []
        gr.Warning(f"Could not load {repo}: {exc}")
        return session.sid, f"**{repo}** failed: {exc}"
    at = f"@{revision}" if revision else ""
    return session.sid, f"**{repo}{at}** - {len(session.rows)} episodes from {note}"


def on_filter(sid, query, sort, current):
    session = get_session(sid)
    if not session.rows:
        return gr.update(choices=[], value=None), []
    rows = filter_rows(session, query, sort)
    choices = [(label_for(r), r["episode_id"]) for r in rows]
    ids = [r["episode_id"] for r in rows]
    value = current if current in ids else (ids[0] if ids else None)
    return (
        gr.update(choices=choices, value=value, label=f"Episode ({len(ids)} of {len(session.rows)})"),
        ids,
    )


def on_select(sid, episode_id, view, show_diff, fps):
    session = get_session(sid)
    if not session.rows or not episode_id:
        return None, "*No episode selected.*", gr.update()
    row = row_by_id(session, episode_id)
    crit = critical_index(row)

    choices = camera_choices(n_cameras(session.pair(episode_id)[0]))
    valid = [v for _, v in choices]
    view = view if view in valid else choices[0][1]

    video = render_video(session, episode_id, view, show_diff, fps, crit)
    return video, metrics_md(row), gr.update(choices=choices, value=view)


def step(ids, current, delta):
    if not ids:
        return None
    try:
        i = ids.index(current)
    except ValueError:
        i = 0
    return ids[(i + delta) % len(ids)]


# --------------------------------------------------------------------------- app

CSS = """
.window-video video { max-height: 82vh; object-fit: contain; background: #111; }
"""

with gr.Blocks(title="Rollout vs. ground truth", fill_width=True, delete_cache=(600, 600)) as demo:
    sid_state = gr.State("", time_to_live=SESSION_TTL, delete_callback=release)
    ids_state = gr.State([])

    with gr.Row():
        with gr.Column(scale=1, min_width=340):
            gr.Markdown("### Rollout vs. ground truth")
            with gr.Row():
                repo_box = gr.Textbox(
                    label="Dataset repo",
                    placeholder="owner/name  (or owner/name@revision, or a dataset URL)",
                    scale=4,
                )
                load_btn = gr.Button("Load", variant="primary", scale=1, min_width=70)
            status = gr.Markdown("Enter a dataset repo id above and press **Load**.")
            query = gr.Textbox(label="Search", placeholder="task or episode id")
            sort = gr.Dropdown(SORTS, value=SORTS[0], label="Sort")
            picker = gr.Dropdown([], label="Episode", filterable=True)
            with gr.Row():
                prev_btn = gr.Button("Prev")
                next_btn = gr.Button("Next")
            view = gr.Radio(camera_choices(3), value=0, label="Camera")
            show_diff = gr.Checkbox(False, label="Add |GT - generated| panel")
            fps = gr.Slider(2, 30, value=8, step=1, label="Playback fps")
            info = gr.Markdown()

        with gr.Column(scale=4):
            video = gr.Video(
                show_label=False,  # the panels are captioned in-frame; the chip covers them
                autoplay=True,
                loop=True,
                buttons=["download", "fullscreen"],
                elem_classes="window-video",
            )

    filter_io = ([sid_state, query, sort, picker], [picker, ids_state])
    select_io = ([sid_state, picker, view, show_diff, fps], [video, info, view])

    # Load and filter only repopulate the episode list; selecting an episode is what renders.
    for event in (load_btn.click, repo_box.submit):
        event(on_load, [repo_box, sid_state], [sid_state, status]).then(on_filter, *filter_io)

    for control in (query, sort):
        control.change(on_filter, *filter_io)

    picker.change(on_select, *select_io)
    for control in (view, show_diff, fps):
        control.change(on_select, *select_io)

    prev_btn.click(lambda i, c: step(i, c, -1), [ids_state, picker], picker)
    next_btn.click(lambda i, c: step(i, c, +1), [ids_state, picker], picker)

if __name__ == "__main__":
    demo.queue(max_size=16).launch(css=CSS)