"""Rerun the AR-view viewer's render on 10 frames from episode 1, with the
TCP site's size cranked up to 3 cm so it shows clearly in every panel.
Answers: is the TCP visible in the wrist camera's FOV or not?
"""
from __future__ import annotations

import os, sys, shutil
from pathlib import Path

os.environ.setdefault("MUJOCO_GL", "osmesa")

DATASET = Path("/data/cameron/puget_ar_datasets/testing_cube_in_bowl")
UMI_SRC = Path("/data/cameron/cad_recovery/mj_umi_v2")
UMI_TMP = Path("/tmp/mj_umi_v2_bigtcp")
OUT = Path("/data/cameron/repos/smithbot_v4/preprocess/kp_diff_viz_bigtcp")
AR_VIEW = Path("/data/cameron/claude_feetech_controller/ar_view")

# 10 frames from ep_auto_mrjt0q63_1 [91..128]
FRAMES = [91, 95, 99, 103, 107, 111, 115, 119, 123, 127]

# Prepare a patched copy of the UMI model with an oversized TCP sphere.
if UMI_TMP.exists():
    shutil.rmtree(UMI_TMP)
shutil.copytree(UMI_SRC, UMI_TMP)
xml_path = UMI_TMP / "umi_gripper_v2.xml"
text = xml_path.read_text()
patched = text.replace(
    '<site name="tcp" pos="-0.1570 0.0248 0.0077" size="0.004" rgba="0.1 0.9 0.2 0.8" group="1" />',
    '<site name="tcp" pos="-0.1570 0.0248 0.0077" size="0.030" rgba="0.1 0.95 0.15 0.9" group="1" />',
)
if patched == text:
    raise RuntimeError("TCP site line not found in umi_gripper_v2.xml (has cad's copy changed?)")
xml_path.write_text(patched)
print(f"patched TCP site size to 0.030 in {xml_path}")

# Also need the site to be in a rendered group. Default MuJoCo renderer shows
# groups 0-5; group=1 already qualifies. Site material/rgba OK. Good.

# Monkey-patch the viewer's UMI_MODEL constant so it uses our patched copy.
sys.path.insert(0, str(AR_VIEW))
import render_dataset_frame as R
R.UMI_MODEL = xml_path
print(f"redirected R.UMI_MODEL → {R.UMI_MODEL}")

# Rebuild argv for the viewer's main().
OUT.mkdir(parents=True, exist_ok=True)
argv_str = [
    "render_dataset_frame.py",
    "--dataset-root", str(DATASET),
    "--frames", ",".join(str(f) for f in FRAMES),
    "--out", str(OUT),
    "--name-fmt", "bigtcp_f{frame:06d}.jpg",
    "--ref", "center",
    "--horizon", "30",
]
sys.argv = argv_str
print("running viewer with argv:", argv_str[1:])
R.main()
print(f"\nrendered → {OUT}/")
print(f"browse: https://omidlab.net/browse/repos/smithbot_v4/preprocess/kp_diff_viz_bigtcp/")