"""Record a localization dataset from the scene + wrist cameras.

World frame = the calib board's ArUco (GridBoard) frame. Everything is solved by PnP on
the boards in the scene image (no MuJoCo needed):

  intrinsics   : scene-camera K, solved once from the first --calib-frames views (saved in meta).
  scene_cam_pose : the scene camera's pose in the world frame, per frame (OpenCV convention,
                   z-forward). If the calib board isn't visible, the last known pose is reused.
  umi_pose     : the UMI gripper's ArUco pose (UMI_V2 board) in the world frame, per frame
                   (reuses last known when not visible; a visibility flag is stored).
  wrist rgb    : the wrist camera frame (downscaled, aspect preserved).
  gripper      : the jaw servo tick + calibrated radians.

Images are saved downscaled (default 500px wide) at --fps (default 4). Layout:
  datasets/<task>/ { meta.json, scene/NNNNNN.jpg, wrist/NNNNNN.jpg, state/NNNNNN.npz }

Usage:
  python record_dataset.py <task> [--scene-index 1] [--wrist-index 0] [--fps 4]
        [--scene-width 500] [--wrist-width 500] [--calib-frames 10] [--overwrite] [--port ...]
"""
import argparse
import json
import os
import sys
import time

import numpy as np
import cv2

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from feetech_bus import Bus, ticks_to_rad                    # noqa: E402
from configs import calib_path                               # noqa: E402
from calib_utils import detect_boards, estimate_focal, solve_pose  # noqa: E402
from calib_board import CALIB_BOARD                           # noqa: E402
from aruco_boards import UMI_V2                               # noqa: E402
from cameras import open_named, load_intrinsics               # noqa: E402

HERE = os.path.dirname(os.path.abspath(__file__))
DATASETS = os.path.join(HERE, "datasets")


def downscale(frame, width):
    h = int(round(frame.shape[0] * width / frame.shape[1]))
    return cv2.resize(frame, (width, h))


def to_height(frame, h):
    return cv2.resize(frame, (int(round(frame.shape[1] * h / frame.shape[0])), h))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("task")
    ap.add_argument("--scene-index", type=int, default=1)
    ap.add_argument("--wrist-index", type=int, default=0)
    ap.add_argument("--fps", type=float, default=4.0)
    ap.add_argument("--scene-width", type=int, default=500)
    ap.add_argument("--wrist-width", type=int, default=500)
    ap.add_argument("--calib-frames", type=int, default=10)
    ap.add_argument("--overwrite", action="store_true")
    ap.add_argument("--no-preview", action="store_true")
    ap.add_argument("--port", default=None)
    a = ap.parse_args()

    # gripper (jaw) servo from calib/gripper.json
    try:
        gcal = json.load(open(calib_path("gripper")))
    except FileNotFoundError:
        sys.exit("No calib/gripper.json. Calibrate the gripper servo first.")
    jaw = next(((int(sid), int(c["sign"])) for sid, c in gcal.items()
                if c["joint"] in ("jaw", "gripper_jaw")), None)
    if jaw is None:
        sys.exit("No jaw/gripper_jaw servo in calib/gripper.json.")
    jaw_id, jaw_sign = jaw

    # dataset dirs
    ddir = os.path.join(DATASETS, a.task)
    sdir, wdir, stdir = (os.path.join(ddir, d) for d in ("scene", "wrist", "state"))
    if os.path.isdir(sdir) and os.listdir(sdir) and not a.overwrite:
        sys.exit(f"Dataset '{a.task}' exists. Use --overwrite or a new name.")
    for d in (sdir, wdir, stdir):
        os.makedirs(d, exist_ok=True)
        for f in os.listdir(d):
            os.remove(os.path.join(d, f))

    # cameras
    scap, sidx, sfr = open_named("scene", a.scene_index)
    wcap, widx, wfr = open_named("wrist", a.wrist_index)
    if scap is None or sfr is None:
        sys.exit("Scene camera didn't open / no frames.")
    if wcap is None:
        sys.exit("Wrist camera didn't open.")
    natW, natH = sfr.shape[1], sfr.shape[0]
    print(f"scene cam idx {sidx} @ {natW}x{natH}   wrist cam idx {widx}")

    # scene intrinsics: solve once from the first N calib-board views (pinhole)
    print(f"Recording the first {a.calib_frames} frames for scene intrinsics (point at the calib board)...")
    views = []
    while len(views) < a.calib_frames:
        ok, frame = scap.read()
        if not ok or frame is None:
            continue
        d = detect_boards(cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY), [CALIB_BOARD])
        if "calib" in d:
            views.append(d["calib"])
        if not a.no_preview:
            cv2.putText(frame, f"scene intrinsics {len(views)}/{a.calib_frames}", (10, 30),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 200, 255), 2)
            cv2.imshow("record_dataset", downscale(frame, 900))
            cv2.waitKey(1)
    K, _, rms = estimate_focal(views, (natW, natH), f_guess=0.9 * max(natW, natH))
    s = a.scene_width / float(natW)
    K_save = K.copy(); K_save[0] *= s; K_save[1] *= s
    print(f"  scene f={K[0,0]:.0f}px @ {natW}x{natH}, reproj {rms:.2f}px")

    # wrist intrinsics: saved (the B0332 lens is fixed; only the iPhone drifts). Scaled to
    # the wrist's current capture resolution so it can localize from the calib board too.
    natWw, natHw = wfr.shape[1], wfr.shape[0]
    wintr = load_intrinsics("wrist")
    if wintr is not None:
        sw = natWw / float(wintr["width"])
        K_wrist = wintr["K"].copy(); K_wrist[0] *= sw; K_wrist[1] *= sw
        sws = a.wrist_width / float(natWw)
        K_wrist_save = K_wrist.copy(); K_wrist_save[0] *= sws; K_wrist_save[1] *= sws
        print(f"  wrist intrinsics loaded: f={K_wrist[0,0]:.0f}px @ {natWw}x{natHw}")
    else:
        K_wrist = K_wrist_save = None
        print("  wrist intrinsics NOT found (run calibrate_camera.py wrist) — wrist pose off, RGB still saved")

    # bus
    bus = Bus(a.port) if a.port else Bus()
    if not bus.ping(jaw_id):
        sys.exit(f"Gripper servo {jaw_id} not responding.")
    bus.torque(jaw_id, False)

    def board_meta(spec):
        return dict(ids=[int(spec["ids"][0]), int(spec["ids"][-1])], shape=list(spec["shape"]),
                    marker_m=spec["marker_m"], sep_m=spec["sep_m"])
    meta = dict(task=a.task, fps=a.fps, jaw_servo_id=jaw_id, jaw_sign=jaw_sign,
                camera_convention="opencv_z_forward", distortion="zero (pinhole)",
                world_frame="calib board GridBoard frame (ids 109-132)",
                boards=dict(calib=board_meta(CALIB_BOARD), umi=board_meta(UMI_V2)),
                scene=dict(native_wh=[natW, natH], save_width=a.scene_width,
                           K_native=K.tolist(), K_save=K_save.tolist(),
                           intrinsics_reproj_px=float(rms), intrinsics_source="live (this run)"),
                wrist=dict(native_wh=[natWw, natHw], save_width=a.wrist_width,
                           K_native=(K_wrist.tolist() if K_wrist is not None else None),
                           K_save=(K_wrist_save.tolist() if K_wrist_save is not None else None),
                           intrinsics_source=("saved intrinsics/wrist.json" if K_wrist is not None else None)),
                started=time.strftime("%Y-%m-%d %H:%M:%S"))
    json.dump(meta, open(os.path.join(ddir, "meta.json"), "w"), indent=2)

    print(f"\nRecording '{a.task}' at {a.fps} fps -> {ddir}\n  'q' in the preview (or Ctrl-C) to stop.\n")
    last_scene = np.eye(4)
    last_umi = np.eye(4)
    last_wrist = np.eye(4)
    period = 1.0 / a.fps
    i = 0
    jpg = [cv2.IMWRITE_JPEG_QUALITY, 95]
    try:
        while True:
            t0 = time.perf_counter()
            oks, sframe = scap.read()
            okw, wframe = wcap.read()
            if not oks or sframe is None:
                break

            dets = detect_boards(cv2.cvtColor(sframe, cv2.COLOR_BGR2GRAY), [CALIB_BOARD, UMI_V2])
            scene_vis = "calib" in dets
            if scene_vis:
                Tcb_cam, _ = solve_pose(*dets["calib"], K)       # calib board -> cam
                last_scene = np.linalg.inv(Tcb_cam)              # camera pose in world
            scene_pose = last_scene
            umi_vis = "umi_v2" in dets
            if umi_vis:
                Tumi_cam, _ = solve_pose(*dets["umi_v2"], K)     # umi board -> cam
                last_umi = scene_pose @ Tumi_cam                 # umi pose in world
            umi_pose = last_umi

            # wrist camera pose: detect the calib (scene) board in the wrist view
            wrist_vis = False
            if okw and wframe is not None and K_wrist is not None:
                wd = detect_boards(cv2.cvtColor(wframe, cv2.COLOR_BGR2GRAY), [CALIB_BOARD])
                if "calib" in wd:
                    Tcb_wcam, _ = solve_pose(*wd["calib"], K_wrist)  # calib board -> wrist cam
                    last_wrist = np.linalg.inv(Tcb_wcam)             # wrist camera pose in world
                    wrist_vis = True
            wrist_pose = last_wrist

            tick = bus.read_pos(jaw_id)
            rad = ticks_to_rad(tick, jaw_sign) if tick is not None else float("nan")

            name = f"{i:06d}"
            cv2.imwrite(os.path.join(sdir, f"{name}.jpg"), downscale(sframe, a.scene_width), jpg)
            if okw and wframe is not None:
                cv2.imwrite(os.path.join(wdir, f"{name}.jpg"), downscale(wframe, a.wrist_width), jpg)
            np.savez(os.path.join(stdir, f"{name}.npz"),
                     scene_cam_pose=scene_pose, umi_pose=umi_pose, wrist_cam_pose=wrist_pose,
                     scene_visible=scene_vis, umi_visible=umi_vis, wrist_visible=wrist_vis,
                     gripper_tick=tick if tick is not None else -1, gripper_rad=rad, t=time.time())
            i += 1

            if not a.no_preview:
                ds = to_height(sframe, 480)
                cv2.putText(ds, f"REC #{i}  scene:{'Y' if scene_vis else 'n'}  umi:{'Y' if umi_vis else 'n'}"
                            f"  grip:{tick}", (10, 28), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
                if okw and wframe is not None:
                    dw = to_height(wframe, 480)
                    cv2.putText(dw, f"wrist  board:{'Y' if wrist_vis else 'n'}", (10, 28),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
                    disp = np.hstack([ds, dw])
                else:
                    disp = ds
                cv2.imshow("record_dataset", disp)
                if cv2.waitKey(1) & 0xFF == ord("q"):
                    break

            dt = time.perf_counter() - t0
            if dt < period:
                time.sleep(period - dt)
    except KeyboardInterrupt:
        print("\n  stopped (Ctrl-C).")
    finally:
        meta.update(num_frames=i, ended=time.strftime("%Y-%m-%d %H:%M:%S"))
        json.dump(meta, open(os.path.join(ddir, "meta.json"), "w"), indent=2)
        scap.release(); wcap.release()
        cv2.destroyAllWindows()
        bus.close()
        print(f"\nSaved {i} frames to {ddir}")


if __name__ == "__main__":
    main()