"""Exo-aruco solve + frame conventions.

Concise: ~5 functions. Used by unit_tests/calibration_test.py (offline)
and calibrate.py (live, via chiral). Aruco failures on individual
frames raise; pool_exo_aruco filters those when collecting multiple.

Conventions:
- ``solve_exo_aruco`` returns ``T_scene_in_right_arm_base``.
- ``to_world(T, T_lfr)`` lifts into world (= left_arm_base) via T_lfr.
- ``T_lfr`` is stored correctly in training PKLs under
  ``T_left_from_right``. The JSON key ``right_base_to_left_base`` is
  misnamed (stores its inverse) — load from a PKL instead.
"""
from __future__ import annotations

import pickle
import sys
import time
from pathlib import Path
from typing import Callable, Optional

import numpy as np


def _exo_modules():
    """Lazy-import the exo_redo helpers (depend on sys.path tweak)."""
    p = str(Path(__file__).resolve().parents[1] / "raiden_fork" /
            "third_party" / "exo_redo")
    if p not in sys.path:
        sys.path.insert(0, p)
    from ExoConfigs.yam_exo import YAM_BASE_ONLY_CONFIG
    from ExoConfigs.exoskeleton import link_to_aruco_transform
    from exo_utils import do_est_aruco_pose, ARUCO_DICT
    return YAM_BASE_ONLY_CONFIG, do_est_aruco_pose, ARUCO_DICT, link_to_aruco_transform


def solve_exo_aruco(scene_bgr: np.ndarray, scene_K: np.ndarray) -> np.ndarray:
    """T_cam_in_right_arm_base from a single frame. Raises on no detection.

    Three steps:
    1. ``do_est_aruco_pose`` → ``T_board_in_cam`` (with raiden's Y/Z flip).
    2. Compose with the fixed ``T_aruco_to_link`` to get ``T_link_in_cam``.
    3. Invert to get ``T_cam_in_link`` (= ``T_cam_in_right_arm_base``).

    ``do_est_aruco_pose`` expects **BGR** input.
    """
    cfg, do_est, ARUCO_DICT, link_to_aruco = _exo_modules()
    link_cfg = cfg.links["larger_coarse_board"]
    board = cfg.aruco_board_objects["larger_coarse_board"]
    r = do_est(
        scene_bgr, ARUCO_DICT, board, link_cfg.board_length,
        cameraMatrix=scene_K.astype(np.float64),
        distCoeffs=np.zeros(5, dtype=np.float64),
    )
    if isinstance(r, int) and r == -1:
        raise RuntimeError("aruco detection failed")
    T_board_in_cam = np.asarray(r["est_aruco_pose"], dtype=np.float64)
    T_aruco_to_link = np.linalg.inv(link_to_aruco(link_cfg))
    T_link_in_cam = T_board_in_cam @ T_aruco_to_link
    return np.linalg.inv(T_link_in_cam)


def load_T_lfr(path: Path) -> np.ndarray:
    """T_left_from_right from either a ``.npy`` (just the 4×4 matrix) or
    a training PKL (whole frame dict with ``T_left_from_right`` key).

    The lab-side canonical is ``cal/T_left_from_right.npy``.
    """
    path = Path(path)
    if path.suffix == ".npy":
        return np.asarray(np.load(path), dtype=np.float64)
    with open(path, "rb") as f:
        fd = pickle.load(f)
    return np.asarray(fd["T_left_from_right"], dtype=np.float64)


def to_world(T_in_rbase: np.ndarray, T_lfr: np.ndarray) -> np.ndarray:
    """Convert right_arm_base → world (= left_arm_base)."""
    return T_lfr @ T_in_rbase


def average_transforms(Ts) -> np.ndarray:
    """Mean translation + SVD-orthogonalized mean rotation."""
    Ts = np.stack(Ts)
    t = Ts[:, :3, 3].mean(0)
    R = Ts[:, :3, :3].mean(0)
    U, _, Vt = np.linalg.svd(R)
    R_orth = U @ Vt
    if np.linalg.det(R_orth) < 0:
        U[:, -1] *= -1
        R_orth = U @ Vt
    out = np.eye(4)
    out[:3, :3] = R_orth
    out[:3, 3] = t
    return out


def pool_exo_aruco(
    get_frame: Callable[[], tuple[np.ndarray, np.ndarray]],
    min_detections: int = 8,
    max_wait_s: Optional[float] = 20.0,
) -> np.ndarray:
    """Pool N exo-aruco detections via a get_frame() callback → (rgb, K).

    Caller supplies the frame source — closure over chiral latest_obs for
    live, or a generator over saved frames for tests. Dedupes by id(rgb).
    Per-frame aruco failures are filtered (not silent fallbacks — we only
    aggregate successful detections; if we never reach N we raise).

    Set ``max_wait_s=None`` to wait indefinitely (useful when the user is
    still physically positioning the board after launch). Prints a
    progress hint every ~5 s while waiting.
    """
    detections, last_id, t0 = [], None, time.time()
    last_progress = t0
    while len(detections) < min_detections:
        now = time.time()
        if max_wait_s is not None and now - t0 > max_wait_s:
            raise TimeoutError(
                f"only {len(detections)} of {min_detections} detections "
                f"in {max_wait_s:.1f}s"
            )
        if now - last_progress > 5.0:
            print(f"  …pooling aruco: {len(detections)}/{min_detections} "
                  f"(elapsed {now - t0:.0f}s)", flush=True)
            last_progress = now
        rgb, K = get_frame()
        cur = id(rgb)
        if cur == last_id:
            time.sleep(0.02)
            continue
        last_id = cur
        try:
            detections.append(solve_exo_aruco(rgb, K))
        except RuntimeError:
            continue
    return average_transforms(detections)