"""Live real-vs-sim AR overlay with two render layers, two named cameras.

Cameras are referenced by role: "scene" (third-person, top row) and "wrist" (UMI wrist
cam, bottom row), resolved via cameras.json. Each uses its saved intrinsics from
calibrate_camera.py (frames are undistorted and the pose solved with the calibrated K);
if a camera has no saved intrinsics yet, it falls back to a live single-view focal estimate.

Grid per frame (posed by the live calibrated servos):
  scene row: [ rgb | scene render | umi render (if UMI board seen) ]
  wrist row: [ rgb | scene render (if base board seen) | blank ]

Usage:
  python arview.py --config arm [--scene-index N] [--wrist-index N] [--width 640] [--port ...]
"""
import argparse
import json
import os
import sys

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                                  # noqa: E402
from configs import CONFIGS, calib_path                      # noqa: E402
from calib_utils import detect_boards, estimate_focal        # noqa: E402
from aruco_boards import ARM_BASE, UMI_V2                     # noqa: E402
from calib_board import CALIB_BOARD                           # noqa: E402
from render_core import Layer, label, blank, blend_cell       # noqa: E402
from cameras import open_named, load_intrinsics               # noqa: E402


def merged_calib(*configs):
    out = {}
    for c in configs:
        try:
            data = json.load(open(calib_path(c)))
        except FileNotFoundError:
            continue
        for sid, e in data.items():
            out.setdefault(sid, e)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default="arm")
    ap.add_argument("--scene-index", type=int, default=1)
    ap.add_argument("--wrist-index", type=int, default=0)
    ap.add_argument("--width", type=int, default=640)
    ap.add_argument("--board", choices=["arm_base", "calib"], default="calib",
                    help="scene-localization board: the large calib board (default) or the small arm-base board")
    ap.add_argument("--live-calib", default="scene",
                    help="comma-sep camera roles to intrinsic-calibrate live at startup instead of "
                         "loading a saved file (default: scene — e.g. the iPhone, whose intrinsics drift)")
    ap.add_argument("--calib-frames", type=int, default=10)
    ap.add_argument("--max-display", type=int, default=1600,
                    help="cap the montage width for display (detection is full-res regardless)")
    ap.add_argument("--port", default=None)
    a = ap.parse_args()
    live_calib = {s.strip() for s in a.live_calib.split(",") if s.strip()}

    # Each layer poses from its OWN config's calib (arm servos vs the gripper servo can share
    # an id, so a single merged file would relabel one of them).
    arm_calib = merged_calib(a.config)
    grip_calib = merged_calib("gripper")
    if not arm_calib and not grip_calib:
        sys.exit("No calibration found. Run calibrate_servo.py first.")

    print("Measuring board frames (loop closure)...")
    if a.board == "calib":     # standalone calib board (not attached to the arm): render the board only
        scene_model = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                   "calib_board_model", "calib_board.xml")
        scene_layer = Layer("scene", scene_model, CALIB_BOARD, "calib", lambda jn: jn)
    else:
        scene_layer = Layer("scene", CONFIGS[a.config]["model"], ARM_BASE, "armboard", lambda jn: jn)
    umi_layer = Layer("umi", CONFIGS["gripper"]["model"], UMI_V2, "ar_back",
                      lambda jn: "jaw" if jn in ("jaw", "gripper_jaw") else None)
    layers = [scene_layer, umi_layer]
    scene_layer.build_posing(arm_calib)
    umi_layer.build_posing(grip_calib)
    for L in layers:
        print(f"  {L.label}: board reproj {L.reproj:.2f}px, posed servos {[s for s, _, _ in L.posing]}")

    bus = Bus(a.port) if a.port else Bus()
    all_sids = sorted({int(s) for s in list(arm_calib) + list(grip_calib)})
    live = [sid for sid in all_sids if bus.ping(sid)]
    for sid in live:
        bus.torque(sid, False)
    print(f"Live servos: {live}")

    cams = {}
    for role, ovr in [("scene", a.scene_index), ("wrist", a.wrist_index)]:
        cap, idx, _ = open_named(role, ovr)
        if cap is None:
            print(f"  {role}: no camera — skipping")
            continue
        intr = None if role in live_calib else load_intrinsics(role)
        cams[role] = {"cap": cap, "idx": idx, "intr": intr, "K": None}
        if role in live_calib:
            print(f"  {role}: index {idx}, will live-calibrate intrinsics at startup ({a.calib_frames} frames)")
        elif intr is None:
            print(f"  {role}: index {idx}, NO saved intrinsics — estimating live "
                  f"(run: python calibrate_camera.py {role} --index {idx})")
        else:
            print(f"  {role}: index {idx}, intrinsics loaded "
                  f"(rms {intr['rms']:.2f}px @ {intr['width']}x{intr['height']})")
    if not cams:
        sys.exit("No cameras opened — check connections / cameras.json / --scene-index/--wrist-index.")
    specs = [L.board_spec for L in layers]
    Hdef = int(round(a.width * 3 / 4))
    cv2.namedWindow("AR view (q to quit)", cv2.WINDOW_AUTOSIZE)

    # Live intrinsic calibration for flagged cameras (e.g. the iPhone scene cam): collect the
    # first N board detections, multi-view pinhole solve, freeze K for the session.
    board_for_role = {"scene": layers[0].board_spec}
    for role in list(cams):
        if role not in live_calib:
            continue
        cam = cams[role]
        spec = board_for_role.get(role, ARM_BASE)
        print(f"\nHold on, recording the first {a.calib_frames} frames for intrinsic calibration "
              f"of the {role} camera (point it at the {spec['name']} board)...")
        views, nat = [], None
        while len(views) < a.calib_frames:
            ok, frame = cam["cap"].read()
            if not ok or frame is None:
                continue
            nat = (frame.shape[1], frame.shape[0])                  # full-res detection
            d = detect_boards(cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY), [spec])
            if spec["name"] in d:
                views.append(d[spec["name"]])
            disp = cv2.resize(frame, (a.width, int(round(frame.shape[0] * a.width / frame.shape[1]))))
            cv2.imshow("AR view (q to quit)",
                       label(disp, f"calibrating {role} @ {nat[0]}x{nat[1]}: {len(views)}/{a.calib_frames}",
                             (0, 200, 255)))
            if cv2.waitKey(1) & 0xFF == 27:
                break
        if len(views) >= 4:
            cam["K"], _, err = estimate_focal(views, nat, f_guess=0.9 * max(nat))
            print(f"  {role} live intrinsics @ {nat[0]}x{nat[1]}: f={cam['K'][0,0]:.0f}px, reproj {err:.2f}px")
        else:
            print(f"  {role}: not enough board views — falling back to single-frame estimate")

    print("  (window opening — if you don't see it, cmd-tab to the python process)")

    try:
        while True:
            ticks = {}
            for sid in live:
                p = bus.read_pos(sid)
                if p is not None:
                    ticks[sid] = p
            for L in layers:
                L.pose(ticks)

            rows = []
            for role, cam in cams.items():
                try:
                    ok, frame = cam["cap"].read()
                except Exception:
                    ok, frame = False, None
                if not ok or frame is None:
                    ph = blank(Hdef, a.width, f"{role}: no frame — check camera {cam['idx']}")
                    rows.append(np.hstack([ph, blank(Hdef, a.width), blank(Hdef, a.width)]))
                    continue

                natW, natH = frame.shape[1], frame.shape[0]
                # Detect + solve pose at FULL capture resolution (accurate); render/display small.
                dets = detect_boards(cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY), specs)
                intr = cam["intr"]
                if intr is not None and cam["K"] is None:       # scale saved K to current capture res (once)
                    s = natW / float(intr["width"])
                    cam["K"] = intr["K"].copy()
                    cam["K"][0] *= s
                    cam["K"][1] *= s
                if intr is None and cam["K"] is None and dets:  # fallback if live-calib didn't run
                    view = max(dets.values(), key=lambda v: len(v[0]))
                    cam["K"], _, _ = estimate_focal([view], (natW, natH), f_guess=0.9 * max(natW, natH))

                W = a.width
                H = int(round(natH * W / natW))
                real = cv2.resize(frame, (W, H))               # display size
                fovy = 2 * np.degrees(np.arctan(natH / 2.0 / cam["K"][1, 1])) if cam["K"] is not None else 60.0

                scene_L, umi_L = layers[0], layers[1]
                cols = [label(real.copy(), f"{role} rgb {natW}x{natH}")]
                cols.append(blend_cell(scene_L, dets, cam["K"], real, fovy))
                cols.append(blend_cell(umi_L, dets, cam["K"], real, fovy) if role == "scene"
                            else blank(H, W, "umi n/a in wrist view"))
                rows.append(np.hstack(cols))

            if rows:
                wmax = max(r.shape[1] for r in rows)
                rows = [cv2.copyMakeBorder(r, 0, 0, 0, wmax - r.shape[1], cv2.BORDER_CONSTANT)
                        for r in rows]
                mont = np.vstack(rows)
                if mont.shape[1] > a.max_display:   # crisp fit to screen (INTER_AREA beats OS shrink)
                    sc = a.max_display / mont.shape[1]
                    mont = cv2.resize(mont, (a.max_display, int(round(mont.shape[0] * sc))),
                                      interpolation=cv2.INTER_AREA)
                cv2.imshow("AR view (q to quit)", mont)
            if cv2.waitKey(1) & 0xFF == ord("q"):
                break
    finally:
        for cam in cams.values():
            cam["cap"].release()
        for L in layers:
            L.close()
        cv2.destroyAllWindows()
        bus.close()


if __name__ == "__main__":
    main()