"""Render a single grid PNG showing the first-frame RGB + sampled query points
(projected from 3D camera-frame back to pixels) for every episode in
3pointr_data/no_actions/stack_can_50s_25f/0001..0075/.

Missing episodes show a black panel labeled with the source video name.
"""
import argparse, json, os
from pathlib import Path

import numpy as np
import cv2
import imageio


OUT_ROOT = "/home/hongyi/PWM/external/3pointr_data/no_actions/stack_can_50s_25f"


def video_label(E: int) -> str:
    if 1 <= E <= 50:
        return f"tableHF_ep{E-1:03d}"
    if 51 <= E <= 75:
        return f"table2_failure_ep{E+34:03d}"
    return f"E={E}"


def project_pts(pts_xyz: np.ndarray, K: np.ndarray) -> np.ndarray:
    """(N,3) cam-frame -> (N,2) pixel coords. Returns float."""
    fx, fy, cx, cy = K[0,0], K[1,1], K[0,2], K[1,2]
    z = np.maximum(pts_xyz[:, 2], 1e-12)
    u = pts_xyz[:, 0] / z * fx + cx
    v = pts_xyz[:, 1] / z * fy + cy
    return np.stack([u, v], axis=1)


def render_episode_panel(E: int, panel_h: int, panel_w: int) -> np.ndarray:
    """Returns (panel_h, panel_w, 3) BGR image with overlay."""
    ep_dir = os.path.join(OUT_ROOT, f"{E:04d}")
    rgb_path = os.path.join(ep_dir, "timestep_1", "1_rgb.png")
    pos_path = os.path.join(ep_dir, "timestep_1", "1_pos.pth")
    vis_path = os.path.join(ep_dir, "timestep_1", "1_visibility.pth")
    meta_path = os.path.join(ep_dir, "metadata.json")

    label = f"{E:02d} {video_label(E)}"

    if not (os.path.exists(rgb_path) and os.path.exists(pos_path) and os.path.exists(meta_path)):
        # MISSING — black panel with label
        panel = np.zeros((panel_h, panel_w, 3), dtype=np.uint8)
        cv2.putText(panel, "MISSING", (8, 30),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2, cv2.LINE_AA)
        cv2.putText(panel, label, (8, panel_h - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255, 255, 255), 1, cv2.LINE_AA)
        return panel

    rgb = imageio.imread(rgb_path)  # (H, W, 3) RGB
    H, W = rgb.shape[:2]
    img = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)

    # Load 3D pos and intrinsics, project, overlay
    import torch
    pts = torch.load(pos_path, map_location="cpu", weights_only=False).numpy()  # (N,3)
    vis = None
    if os.path.exists(vis_path):
        vis = torch.load(vis_path, map_location="cpu", weights_only=False).numpy().astype(bool)
    meta = json.load(open(meta_path))
    cam = next(iter(meta["intrinsics"].values()))
    K = np.array([[cam["fx"], 0, cam["cx"]],
                  [0, cam["fy"], cam["cy"]],
                  [0, 0, 1]], dtype=np.float32)

    uv = project_pts(pts, K)  # (N,2)
    valid = (uv[:,0] >= 0) & (uv[:,0] < W) & (uv[:,1] >= 0) & (uv[:,1] < H) & (pts[:,2] > 1e-6)
    if vis is not None:
        valid_vis = valid & vis
        valid_invis = valid & (~vis)
    else:
        valid_vis = valid
        valid_invis = np.zeros_like(valid)

    # Draw invisible points (gray) first, then visible (red) on top
    for u, v in uv[valid_invis].astype(int):
        cv2.circle(img, (u, v), 1, (128, 128, 128), -1, cv2.LINE_AA)
    for u, v in uv[valid_vis].astype(int):
        cv2.circle(img, (u, v), 1, (0, 0, 255), -1, cv2.LINE_AA)

    # Resize to panel size
    panel = cv2.resize(img, (panel_w, panel_h), interpolation=cv2.INTER_AREA)

    n_vis = int(valid_vis.sum()); n_total = pts.shape[0]
    cv2.putText(panel, label, (4, panel_h - 22),
                cv2.FONT_HERSHEY_SIMPLEX, 0.40, (0, 0, 0), 3, cv2.LINE_AA)
    cv2.putText(panel, label, (4, panel_h - 22),
                cv2.FONT_HERSHEY_SIMPLEX, 0.40, (255, 255, 255), 1, cv2.LINE_AA)
    cv2.putText(panel, f"vis {n_vis}/{n_total}", (4, panel_h - 6),
                cv2.FONT_HERSHEY_SIMPLEX, 0.40, (0, 0, 0), 3, cv2.LINE_AA)
    cv2.putText(panel, f"vis {n_vis}/{n_total}", (4, panel_h - 6),
                cv2.FONT_HERSHEY_SIMPLEX, 0.40, (255, 255, 255), 1, cv2.LINE_AA)
    return panel


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--cols", type=int, default=10)
    ap.add_argument("--panel_h", type=int, default=192)
    ap.add_argument("--panel_w", type=int, default=192)
    ap.add_argument("--out", default="/home/hongyi/PWM/external/3pointr_data/stack_can_50s_25f_first_frames.png")
    args = ap.parse_args()

    cols = int(args.cols)
    rows = int(np.ceil(75 / cols))
    grid = np.zeros((rows * args.panel_h, cols * args.panel_w, 3), dtype=np.uint8)

    for E in range(1, 76):
        r = (E - 1) // cols
        c = (E - 1) % cols
        panel = render_episode_panel(E, args.panel_h, args.panel_w)
        grid[r*args.panel_h:(r+1)*args.panel_h, c*args.panel_w:(c+1)*args.panel_w] = panel

    cv2.imwrite(args.out, grid)
    print(f"wrote {args.out}  grid={rows}x{cols}  size={grid.shape[1]}x{grid.shape[0]}")


if __name__ == "__main__":
    main()