"""Concatenate scene/NNNNNN.jpg + wrist/NNNNNN.jpg → scene_wrist/NNNNNN.jpg.

Also emits `gripper.json` = per-frame gripper_rad + visibility flags, for the
data viewer's client-side auto-episode detection.

Horizontal concat, wrist resized so its height matches scene, JPEG quality 95.
Top-right black-box HUD reads `gripper: <rad> rad` from state/NNNNNN.npz.

Idempotent: re-bakes only when a source (scene/wrist/state) is newer than the composite.
Always re-writes gripper.json.
"""
import sys, math, json, pathlib
import numpy as np
from PIL import Image, ImageDraw, ImageFont

FONT = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 14)
PAD = 4

def _load_state(state_path):
    if not state_path.exists():
        return None
    return np.load(state_path)

def _hud_text(z):
    if z is None or "gripper_rad" not in z.files:
        return "gripper: n/a"
    rad = float(z["gripper_rad"])
    tick = int(z["gripper_tick"]) if "gripper_tick" in z.files else -1
    if tick == -1 or not math.isfinite(rad):
        return "gripper: n/a"
    return f"gripper: {rad:+.3f} rad"

def _draw_hud(canvas, text):
    d = ImageDraw.Draw(canvas)
    l, t, r, b = d.textbbox((0, 0), text, font=FONT)
    tw, th = r - l, b - t
    x1 = canvas.width - tw - 2 * PAD
    y1 = 0
    x2 = canvas.width
    y2 = th + 2 * PAD
    d.rectangle([x1, y1, x2, y2], fill=(0, 0, 0))
    d.text((x1 + PAD - l, y1 + PAD - t), text, fill=(255, 255, 255), font=FONT)

def bake(root):
    scene_dir = root / "scene"
    wrist_dir = root / "wrist"
    state_dir = root / "state"
    out_dir   = root / "scene_wrist"
    out_dir.mkdir(exist_ok=True)
    scene_files = sorted(scene_dir.glob("*.jpg"))
    baked = skipped = missing = 0
    grippers = []      # gripper_rad or null (per frame)
    umi_visible = []   # bool per frame (false if state missing)
    for sp in scene_files:
        name = sp.name
        wp = wrist_dir / name
        stp = state_dir / (sp.stem + ".npz")
        op = out_dir / name
        z = _load_state(stp)
        # gripper series (always)
        if z is not None and "gripper_rad" in z.files and "gripper_tick" in z.files:
            tick = int(z["gripper_tick"]); rad = float(z["gripper_rad"])
            grippers.append(None if (tick == -1 or not math.isfinite(rad)) else rad)
        else:
            grippers.append(None)
        umi_visible.append(bool(z["umi_visible"]) if z is not None and "umi_visible" in z.files else False)
        # composite
        if not wp.exists():
            missing += 1
            continue
        src_mtime = max(sp.stat().st_mtime, wp.stat().st_mtime,
                        stp.stat().st_mtime if stp.exists() else 0)
        if op.exists() and op.stat().st_mtime >= src_mtime:
            skipped += 1
            continue
        s = Image.open(sp)
        w = Image.open(wp)
        if w.height != s.height:
            new_w = round(w.width * s.height / w.height)
            w = w.resize((new_w, s.height), Image.LANCZOS)
        canvas = Image.new("RGB", (s.width + w.width, s.height), (0, 0, 0))
        canvas.paste(s, (0, 0))
        canvas.paste(w, (s.width, 0))
        _draw_hud(canvas, _hud_text(z))
        canvas.save(op, "JPEG", quality=95)
        baked += 1
    # gripper.json sidecar (in dataset root — sibling of scene_wrist/)
    (root / "gripper.json").write_text(json.dumps({
        "num_frames": len(scene_files),
        "gripper_rad": grippers,
        "umi_visible": umi_visible,
    }))
    print(f"{root.name}: baked={baked} skipped={skipped} missing={missing} total={len(scene_files)} gripper.json written")

for name in sys.argv[1:]:
    bake(pathlib.Path.home() / "ar_datasets" / name)