"""Pre-render mesh+hand overlays for all TACO sequences.

Uses the same rendering approach as notebooks/explore_taco.ipynb (cell 18).

Usage:
    .venv/bin/python scripts/prerender_overlays.py

Outputs: outputs/overlays/<sequence_id_sanitized>.jpg
"""

import sys
import numpy as np
import pandas as pd
import cv2
import torch
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

import decord
decord.bridge.set_bridge("torch")

from fwomo_3d.config import TacoDatasetCfg, ViewSamplerAllCfg
from fwomo_3d.data.taco_utils import load_object_meshes, render_mesh_on_image
from fwomo_3d.utils.visualization.hand import (
    HAND_COLORS,
    c2w_and_norm_K_to_w2c,
    draw_hand_skeleton,
)

TACO = Path("/mnt/fast/mzhobro/taco_dataset_resized")
OUT = PROJECT_ROOT / "outputs" / "overlays"
OUT.mkdir(parents=True, exist_ok=True)

GRID_ROWS, GRID_COLS = 3, 4
CAM_W, CAM_H = 384, 282
HEADER_H = 36


def sanitize(seq_id):
    return seq_id.replace("/", "__").replace(" ", "_").replace("(", "").replace(")", "").replace(",", "")


def main():
    # Build dataset (same config as explore_taco notebook)
    cfg = TacoDatasetCfg(
        root=str(TACO),
        image_shape=[376, 512],
        view_sampler=ViewSamplerAllCfg(name="all"),
        near=0.1,
        far=5.0,
        load_hand_poses=True,
        load_object_poses=True,
    )
    # Build all splits to cover every sequence
    ds_train = cfg.build("train")
    ds_val = cfg.build("val")
    ds_test = cfg.build("test")
    # Merge metadata from all splits
    ds = ds_train
    all_meta = pd.concat([ds_train.meta, ds_val.meta, ds_test.meta], ignore_index=True)
    ds.meta = all_meta
    models_dir = TACO / "object_models_released"

    n = len(ds.meta)
    existing = len(list(OUT.glob("*.jpg")))
    print(f"Rendering {n} sequences ({existing} already done)")

    for seq_idx in range(n):
        row = ds.meta.iloc[seq_idx]
        seq_id = row["sequence_id"]
        out_path = OUT / f"{sanitize(seq_id)}.jpg"
        if out_path.exists():
            continue

        try:
            frame_idx = 0
            views = ds.load_sequence_views(seq_idx, frame_idx)
            H, W = views["image"].shape[2], views["image"].shape[3]
            n_cams = len(views["index"])

            # Load meshes
            meshes = load_object_meshes(TACO, row, frame_idx, models_dir)

            # Load hands from actions
            hands = {}
            if "actions" in views:
                for key in ("left_hand", "right_hand"):
                    if key in views["actions"]:
                        joints = views["actions"][key][0].numpy()  # (21, 3)
                        if not np.allclose(joints, 0.0):
                            side = key.replace("_hand", "")
                            hands[side] = joints.astype(np.float64)

            # Build grid
            canvas = np.zeros((GRID_ROWS * CAM_H + HEADER_H, GRID_COLS * CAM_W, 3), dtype=np.uint8)

            for v in range(min(n_cams, GRID_ROWS * GRID_COLS)):
                img_np = (views["image"][v].permute(1, 2, 0).numpy() * 255).astype(np.uint8).copy()
                c2w = views["extrinsics"][v].numpy()
                K_norm = views["intrinsics"][v].numpy()
                R, T, K_px = c2w_and_norm_K_to_w2c(c2w, K_norm, H, W)

                # Render meshes
                for verts_w, faces_arr, color, label in meshes:
                    img_np = render_mesh_on_image(img_np, verts_w, faces_arr, K_px, R, T, color=color)

                # Render hands
                for side, joints in hands.items():
                    img_np = draw_hand_skeleton(img_np, joints, K_px, R, T, HAND_COLORS[side])

                # Camera label
                cam_label = views["index"][v] if isinstance(views["index"][v], str) else str(views["index"][v])
                cv2.putText(img_np, cam_label[-4:], (5, 22), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)

                # Place in grid
                r, c = v // GRID_COLS, v % GRID_COLS
                y0, x0 = r * CAM_H + HEADER_H, c * CAM_W
                resized = cv2.resize(img_np, (CAM_W, CAM_H))
                canvas[y0:y0 + CAM_H, x0:x0 + CAM_W] = resized

            # Header
            cv2.rectangle(canvas, (0, 0), (canvas.shape[1], HEADER_H - 2), (40, 40, 40), -1)
            header = f"{seq_id}  (date={row.get('date', '?')})"
            cv2.putText(canvas, header, (10, 26), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1)

            # Save as BGR
            cv2.imwrite(str(out_path), cv2.cvtColor(canvas, cv2.COLOR_RGB2BGR),
                        [cv2.IMWRITE_JPEG_QUALITY, 85])

        except Exception as e:
            print(f"  ERROR {seq_id}: {e}")

        if (seq_idx + 1) % 50 == 0:
            done = len(list(OUT.glob("*.jpg")))
            print(f"  {seq_idx + 1}/{n} done ({done} total)")

    print(f"Done. {len(list(OUT.glob('*.jpg')))} overlays in {OUT}")


if __name__ == "__main__":
    main()