"""Offline episode visualizer — robot-mask overlay + future-EE traces.

For one recorded episode (SVO2s + robot_data.npz + calibration_results.json),
render per camera:
  - the mujoco arm+board mask composited onto every RGB frame (same render
    stack as the record_testing live panels: raiden.calibration.exo_auto_viz)
  - the end-effector trace for the NEXT 8 output steps, projected into the
    camera (right arm = red, left arm = blue, shrinking dots into the future)

Static cams (scene/ego) use the episode's own auto-cal extrinsics with the
per-camera arm-2 anchor. Wrist cams recompute their extrinsic EVERY frame:
T_cam_in_world = T_arm_in_world @ FK(q_t) @ T_ee_cam (bimanual wrist hand-eye).

Output: /tmp/viz_out/<cam>.mp4 (H.264) — encode + HTML viewer handled by the
caller.

Run on russet:
  cd /home/robot-lab/raiden && MUJOCO_GL=egl PYOPENGL_PLATFORM=egl \
    .venv/bin/python /tmp/viz_episode.py
"""
import json
import time
from pathlib import Path

import cv2
import numpy as np

# Sets MUJOCO_GL + exo_redo sys.path at import time.
from raiden.calibration.exo_auto_viz import (
    _build_bimanual_scene, _render_with_intrinsics, _compose_overlay,
    _discover_arm_qpos, _fk_base_to_ee, _load_wrist_hand_eye)

EP = Path("/home/robot-lab/raiden/data/raw/verify_cameron_calibration_upload/0001")
OUT = Path("/tmp/viz_out")
STRIDE = 3            # 30 fps SVO -> 10 fps output
FPS_OUT = 30.0 / STRIDE
N_FUTURE = 8          # EE trace: next 8 OUTPUT steps (8 * stride/30 s ahead)
ALPHA = 0.55          # mask blend (lighter than cal panels so RGB stays visible)

OUT.mkdir(exist_ok=True)

# ---------------------------------------------------------------- robot data
npz = np.load(EP / "robot_data.npz")
ts_robot = npz["timestamps"]                       # ns epoch
q7_r = npz["follower_r_joint_pos_7d"].astype(np.float64)
q7_l = npz["follower_l_joint_pos_7d"].astype(np.float64)
n_rob = len(ts_robot)
print(f"[viz] robot samples: {n_rob} over {(ts_robot[-1]-ts_robot[0])/1e9:.1f}s")

# Precompute EE position in each ARM BASE frame for every robot sample.
t0 = time.monotonic()
p_ee_r = np.empty((n_rob, 3)); p_ee_l = np.empty((n_rob, 3))
for i in range(n_rob):
    p_ee_r[i] = _fk_base_to_ee(q7_r[i])[:3, 3]
    p_ee_l[i] = _fk_base_to_ee(q7_l[i])[:3, 3]
print(f"[viz] FK precompute: {time.monotonic()-t0:.1f}s")


def nearest_idx(t_ns):
    i = np.searchsorted(ts_robot, t_ns)
    if i <= 0:
        return 0
    if i >= n_rob:
        return n_rob - 1
    return i if ts_robot[i] - t_ns < t_ns - ts_robot[i - 1] else i - 1


# ---------------------------------------------------------------- calibration
cal = json.loads((EP / "calibration_results.json").read_text())
cams_cal = cal["cameras"]
bt = cal.get("bimanual_transform", {})
t = bt.get("translation_m_right_base_in_left_base")
margin = float(abs(t[1])) if isinstance(t, list) and len(t) == 3 else 0.591
T_arm1_in_world = np.eye(4); T_arm1_in_world[1, 3] = -margin / 2.0
T_arm2_in_world_ideal = np.eye(4); T_arm2_in_world_ideal[1, 3] = +margin / 2.0
print(f"[viz] margin: {margin:.4f} m")

wrist_he = _load_wrist_hand_eye(
    Path.home() / ".config" / "raiden" / "wrist_calibration_result_bimanual.json")
print(f"[viz] wrist hand-eye: {list(wrist_he)}")

# ---------------------------------------------------------------- mujoco scene
import mujoco
from scipy.spatial.transform import Rotation as R
from exo_utils import position_exoskeleton_meshes, get_link_poses_from_robot

model, data, exo_a1, exo_a2, arm2_body_id = _build_bimanual_scene(
    margin=margin, patch_russet_v2_reverse=True)
arm1_qpos, arm1_grip = _discover_arm_qpos(model, prefix="")
arm2_qpos, arm2_grip = _discover_arm_qpos(model, prefix="arm2_")


def drive_qpos(i_rob):
    for idx, qv in zip(arm1_qpos, q7_r[i_rob][:6]):
        data.qpos[idx] = qv
    for qidx, sign, stroke in arm1_grip:
        data.qpos[qidx] = sign * np.clip(q7_r[i_rob][6], 0, 1) * stroke
    for idx, qv in zip(arm2_qpos, q7_l[i_rob][:6]):
        data.qpos[idx] = qv
    for qidx, sign, stroke in arm2_grip:
        data.qpos[qidx] = sign * np.clip(q7_l[i_rob][6], 0, 1) * stroke


def arm2_anchor_from_pair(arm1_key, arm2_key):
    """Per-camera arm-2 world anchor from that cam's own PnP pair."""
    if arm1_key not in cams_cal or arm2_key not in cams_cal:
        return None
    T_c_a1 = np.asarray(cams_cal[arm1_key]["extrinsics"]["T_cam_in_base"])
    T_c_a2 = np.asarray(cams_cal[arm2_key]["extrinsics"]["T_cam_in_base"])
    return T_arm1_in_world @ (T_c_a1 @ np.linalg.inv(T_c_a2))


def apply_arm2_anchor(T_arm2_world):
    model.body_pos[arm2_body_id] = T_arm2_world[:3, 3]
    q_xyzw = R.from_matrix(T_arm2_world[:3, :3]).as_quat()
    model.body_quat[arm2_body_id] = q_xyzw[[3, 0, 1, 2]]


# Canonical arm-2 anchor (scene cam pair) for views without their own pair.
T_arm2_scene = arm2_anchor_from_pair("scene_camera__arm1", "scene_camera__arm2")
if T_arm2_scene is None:
    T_arm2_scene = T_arm2_in_world_ideal

# ------------------------------------------------------- per-camera configs
CAM_SPECS = {}  # name -> dict(static T_world_in_cam OR wrist fn, T_arm2_world)
for name in ("scene_camera", "ego_camera"):
    a1k, a2k = f"{name}__arm1", f"{name}__arm2"
    if a1k not in cams_cal:
        print(f"[viz] {name}: no __arm1 cal — skipped")
        continue
    T_cam_in_world = T_arm1_in_world @ np.asarray(
        cams_cal[a1k]["extrinsics"]["T_cam_in_base"])
    anchor = arm2_anchor_from_pair(a1k, a2k)
    CAM_SPECS[name] = {
        "static_T_world_in_cam": np.linalg.inv(T_cam_in_world),
        "T_arm2_world": anchor if anchor is not None else T_arm2_scene,
    }
for name, (arm_label, T_ee_cam) in wrist_he.items():
    T_arm = T_arm1_in_world if arm_label == "right" else T_arm2_scene
    CAM_SPECS[name] = {
        "wrist_arm": arm_label, "T_ee_cam": T_ee_cam, "T_arm_in_world": T_arm,
        "T_arm2_world": T_arm2_scene,
    }
print(f"[viz] cameras: {list(CAM_SPECS)}")

# ------------------------------------------------------------ trace drawing
TRACE = (("R", p_ee_r, (0, 0, 255)),      # right arm: red (BGR)
         ("L", p_ee_l, (255, 128, 0)))    # left arm: blue


def draw_ee_traces(img, K, T_world_in_cam, i_rob, ts_frame_ns, T_arm2_world):
    """Project current + next-N_FUTURE-step EE world points into the view."""
    dt_ns = int(STRIDE / 30.0 * 1e9)
    idxs = [i_rob] + [nearest_idx(ts_frame_ns + k * dt_ns)
                      for k in range(1, N_FUTURE + 1)]
    for label, p_base, color in TRACE:
        T_arm = T_arm1_in_world if label == "R" else T_arm2_world
        pts = []
        for i in idxs:
            pw = T_arm @ np.array([*p_base[i], 1.0])
            pc = T_world_in_cam @ pw
            if pc[2] < 0.05:
                pts.append(None)
                continue
            uv = K @ (pc[:3] / pc[2])
            pts.append((int(round(uv[0])), int(round(uv[1]))))
        for a, b in zip(pts[:-1], pts[1:]):
            if a is not None and b is not None:
                cv2.line(img, a, b, color, 2, cv2.LINE_AA)
        for k, pt in enumerate(pts):
            if pt is None:
                continue
            r = max(2, 8 - k)
            cv2.circle(img, pt, r, color, -1 if k == 0 else 2, cv2.LINE_AA)
        if pts[0] is not None:
            cv2.putText(img, label, (pts[0][0] + 10, pts[0][1] - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2, cv2.LINE_AA)


# ------------------------------------------------------------- per-cam render
import pyzed.sl as sl

renderer = None
renderer_shape = None

for name, spec in CAM_SPECS.items():
    svo = EP / "cameras" / f"{name}.svo2"
    ip = sl.InitParameters()
    ip.set_from_svo_file(str(svo))
    ip.depth_mode = sl.DEPTH_MODE.NONE
    ip.svo_real_time_mode = False
    zed = sl.Camera()
    if zed.open(ip) != sl.ERROR_CODE.SUCCESS:
        print(f"[viz] {name}: SVO open failed — skipped")
        continue
    info = zed.get_camera_information()
    lc = info.camera_configuration.calibration_parameters.left_cam
    K = np.array([[lc.fx, 0, lc.cx], [0, lc.fy, lc.cy], [0, 0, 1.0]])
    W = info.camera_configuration.resolution.width
    H = info.camera_configuration.resolution.height
    n_frames = zed.get_svo_number_of_frames()
    print(f"[viz] {name}: {W}x{H}, {n_frames} frames, stride {STRIDE}")

    if renderer_shape != (W, H):
        renderer = mujoco.Renderer(model, height=H, width=W)
        renderer_shape = (W, H)

    # Static cams: fixed extrinsic + fixed arm-2 anchor for the whole video.
    apply_arm2_anchor(spec["T_arm2_world"])
    mujoco.mj_forward(model, data)
    position_exoskeleton_meshes(
        exo_a2, model, data, get_link_poses_from_robot(exo_a2, model, data))

    vw = cv2.VideoWriter(str(OUT / f"{name}_raw.mp4"),
                         cv2.VideoWriter_fourcc(*"mp4v"), FPS_OUT, (W, H))
    mat = sl.Mat()
    fidx = 0
    written = 0
    t_cam0 = time.monotonic()
    while True:
        rc = zed.grab()
        if rc != sl.ERROR_CODE.SUCCESS:
            break
        if fidx % STRIDE:
            fidx += 1
            continue
        fidx += 1
        zed.retrieve_image(mat, sl.VIEW.LEFT)
        bgr = mat.get_data()[:, :, :3].copy()
        ts = zed.get_timestamp(sl.TIME_REFERENCE.IMAGE).get_nanoseconds()
        i_rob = nearest_idx(ts)

        drive_qpos(i_rob)
        if "wrist_arm" in spec:
            q7 = q7_r[i_rob] if spec["wrist_arm"] == "right" else q7_l[i_rob]
            T_cam_in_world = (spec["T_arm_in_world"] @ _fk_base_to_ee(q7)
                              @ spec["T_ee_cam"])
            T_world_in_cam = np.linalg.inv(T_cam_in_world)
        else:
            T_world_in_cam = spec["static_T_world_in_cam"]
        mujoco.mj_forward(model, data)

        rendered_rgb, seg = _render_with_intrinsics(
            model, data, T_world_in_cam, K, W, H, renderer=renderer)
        overlay = _compose_overlay(bgr, rendered_rgb, seg_mask=seg,
                                   alpha=ALPHA)
        frame = cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR)
        fg = (seg[:, :, 0] != -1).astype(np.uint8) * 255
        contours, _ = cv2.findContours(fg, cv2.RETR_EXTERNAL,
                                       cv2.CHAIN_APPROX_SIMPLE)
        cv2.drawContours(frame, contours, -1, (0, 0, 255), 1, cv2.LINE_AA)

        draw_ee_traces(frame, K, T_world_in_cam, i_rob, ts,
                       spec["T_arm2_world"])
        cv2.putText(frame, f"{name}  t={(ts-ts_robot[0])/1e9:6.2f}s",
                    (12, 28), cv2.FONT_HERSHEY_SIMPLEX, 0.8,
                    (255, 255, 255), 2, cv2.LINE_AA)
        vw.write(frame)
        written += 1
        if written % 100 == 0:
            print(f"[viz] {name}: {written} frames "
                  f"({(time.monotonic()-t_cam0)/written*1000:.0f} ms/frame)")
    vw.release()
    zed.close()
    print(f"[viz] {name}: DONE — {written} frames in "
          f"{time.monotonic()-t_cam0:.0f}s -> {OUT}/{name}_raw.mp4")

print("[viz] ALL DONE")