#!/usr/bin/env python3
"""仿真相机的几何验收：把深度反投影到 world，看它是否真的落在该落的平面上。

=============================================================================
为什么需要这个工具（而不是"看图有图像就行"）
=============================================================================
仿真里相机的错误几乎都是**静默**的：话题在、图像有、深度有、TF 也查得到，
而三维点位整体是错的。下面每一类都真实踩过：

  1. 深度单位错了（gz 发 32FC1 米，应用当 16UC1 毫米读）→ 差 1000 倍。
  2. frame_id 盖成了相机**本体**系而不是光学系 → tf2 查得到、不报错，
     但反投影出的点整体转错。
  3. **gz 发的 camera_info 是错的**（恒为 fx=fy=277, cx=160, cy=120，
     即它的默认 60° 视场 + 半分辨率主点），而渲染用的是配置的 87° 视场。
     照抄它反投影，点会偏半米量级。
  4. 相机位姿的 rpy 写错 → 相机**背对**目标，只看得见地面。
     这种情况几何仍然自洽（点确实落在真实地面上），所以只查数值发现不了 ——
     必须比对"该看见什么"，也就是本工具要说的平面。

所以本工具的判据是：**反投影出的点必须命中预期平面**。
"有图像"不算通过，"落在正确的位置"才算。

用法（必须先 source scripts/sim_env.sh，因为需要 ROS 环境）：
    python3 scripts/check_sim_cameras.py
    python3 scripts/check_sim_cameras.py --expect d435:0.40:30 --cameras d435

退出码 0 = 全部通过；非 0 = 有不通过的相机（可直接用于 CI / 启动自检）。
=============================================================================
"""

from __future__ import annotations

import argparse
import math
import sys
import time

import numpy as np

try:
    import rclpy
    from rclpy.node import Node
    from rclpy.qos import qos_profile_sensor_data
    from sensor_msgs.msg import CameraInfo, Image
    from tf2_ros import Buffer, TransformListener
except ImportError as exc:  # pragma: no cover - 环境问题，不是逻辑问题
    print(
        f"✘ 导入 ROS 失败：{exc}\n"
        f"  这个脚本要在仿真环境里跑，先执行：\n"
        f"      source scripts/sim_env.sh && . install_sim/setup.bash",
        file=sys.stderr,
    )
    sys.exit(2)


# 默认期望：哪台相机"必须"看见哪个平面。
#   sim_d435 是手外相机，它的职责就是看见台面（P6 的粗定位要靠它），
#        所以要求 ≥10% 的像素落在台面上。
#   sim_d405 是腕上相机（仿真那台 = 左臂），视野窄且会拍到机械臂自己，台面占比放宽到 5%。
#   地面(z=0) 不作为要求：它被拍到多少取决于姿态，不构成"对不对"的判据。
DEFAULT_EXPECT = ("sim_d435:0.40:10", "sim_d405:0.40:5")


def quat_to_R(q) -> np.ndarray:
    """四元数 (x, y, z, w) → 3x3 旋转矩阵。"""
    x, y, z, w = q
    n = x * x + y * y + z * z + w * w
    if n <= 0.0:
        return np.eye(3)
    s = 2.0 / n
    return np.array(
        [
            [1 - s * (y * y + z * z), s * (x * y - z * w), s * (x * z + y * w)],
            [s * (x * y + z * w), 1 - s * (x * x + z * z), s * (y * z - x * w)],
            [s * (x * z - y * w), s * (y * z + x * w), 1 - s * (x * x + y * y)],
        ]
    )


class Grab(Node):
    """抓一帧深度 + 一份内参。"""

    def __init__(self, ns: str) -> None:
        super().__init__(f"check_sim_cameras_{ns}")
        self.ns = ns
        self.img = None
        self.info = None
        self.buf = Buffer()
        self.tl = TransformListener(self.buf, self)
        # QoS 必须与发布端一致：适配器与真机都用 sensor data，
        # 用默认 RELIABLE 订阅会一条都收不到（表现是"话题在但 echo 为空"）。
        self.create_subscription(
            Image,
            f"/{ns}/aligned_depth_to_color/image_raw",
            self._on_img,
            qos_profile_sensor_data,
        )
        self.create_subscription(
            CameraInfo, f"/{ns}/color/camera_info", self._on_info, qos_profile_sensor_data
        )

    def _on_img(self, m) -> None:
        self.img = m

    def _on_info(self, m) -> None:
        self.info = m


def spin_until(node: Node, pred, timeout: float) -> bool:
    t0 = time.time()
    while rclpy.ok() and time.time() - t0 < timeout:
        rclpy.spin_once(node, timeout_sec=0.3)
        if pred():
            return True
    return False


def _tf_rot(node: "Grab", parent: str, child: str) -> np.ndarray:
    """TF parent ← child 的旋转部分（child 的三根轴在 parent 系里的表示）。"""
    tf = node.buf.lookup_transform(parent, child, rclpy.time.Time())
    r = tf.transform.rotation
    return quat_to_R((r.x, r.y, r.z, r.w))


def mount_problems(node: "Grab", ns: str, frame: str, problems: list[str]) -> None:
    """核对**安装朝向**：腕相机视轴要沿法兰轴线，手外相机要朝下看台面。

    为什么单独查这一条（光看"命中平面"不够）：
    腕相机要是朝侧面看，它照样能反投影出 z≈0 的地面 —— 几何自洽、数值"正常"，
    只是**看不见工件**。表现全在应用侧："目标不在相机前方"、伺服一直找工件。
    所以这里查的不是"看见了什么"，而是"朝哪看"：前者会被姿态带偏，后者不会。
    """
    try:
        from arm_sim_gazebo.sim_cameras import (
            MOUNT_TOL_DEG,
            WORLD_LOOK_DOWN_COS,
            flange_frame_for,
            load_cameras,
            mount_axis_deg,
            world_look_down_cos,
        )
    except ImportError as exc:
        # 不静默跳过：安装朝向是本工具**最重要**的一条判据（它拦的是"相机装歪
        # 但数值全正常"那类问题），悄悄不查等于把这条守门人撤了。
        problems.append(
            f"拿不到 sim_cameras 的安装朝向工具（{exc}）—— install_sim 是旧的，"
            f"先跑 bash scripts/build_sim.sh"
        )
        return
    spec = next((s for s in load_cameras() if s.ros_ns == ns), None)
    if spec is None:
        problems.append(f"cameras.yaml 里找不到 ros_ns={ns} 的相机声明，无法核对安装朝向")
        return
    now = rclpy.time.Time()
    flange = flange_frame_for(spec)
    if flange is not None and node.buf.can_transform(flange, frame, now):
        R = _tf_rot(node, flange, frame)
        ang = mount_axis_deg(R)
        print(f"        相对法兰 {flange} 的视轴夹角: {ang:.1f}°  （0° = 正装在法兰轴线上）")
        if ang > MOUNT_TOL_DEG:
            problems.append(
                f"腕部相机视轴与法兰 +Z 差 {ang:.1f}°（要求 ≤{MOUNT_TOL_DEG:.0f}°）—— "
                f"相机装歪了，法兰垂直朝下时它看不到下面的工件。\n"
                f"      90° 左右就是「少转一个直角」：话题/图像/内参全都正常，\n"
                f"      而且反投影仍会命中地面（看的是远处），所以只查平面发现不了。\n"
                f"      修法：devices/sim/gazebo/urdf/arm.sim.sensors.xacro 里"
                f" d405_pitch（正装 = -1.5708）。"
            )
    else:
        cos = world_look_down_cos(_tf_rot(node, "world", frame))
        print(f"        视轴在世界系里的 z 分量: {cos:+.2f}  （-1 = 正朝下看）")
        if cos > WORLD_LOOK_DOWN_COS:
            problems.append(
                f"手外相机没在朝下看台面（视轴 z 分量 {cos:+.2f}，要求 ≤{WORLD_LOOK_DOWN_COS:+.1f}）—— "
                f"位姿 rpy 写错时几何自洽、不报错，只是拍不到工件。\n"
                f"      修法：devices/sim/gazebo/worlds/assembly.world 里该 <model> 的 <pose> pitch。"
            )


def looking_at_table(node: "Grab", frame: str) -> "bool | None":
    """相机此刻是不是**朝下看着台面**（光心在台面 xy 内 + 视轴朝下）。

    为什么要问这一句：腕相机的职责是"悬在工件上方时看清螺纹孔"。机械臂停在初始位
    或别处时它当然只看得见地面（初始位那副姿态里法兰甚至是横着的）—— 这时候拿
    "命中台面 ≥5%"判它不合格是**假警报**，而假警报多了真问题就没人看了。
    拿不到场景配置时返回 None（按"不知道"处理，照常判）。
    """
    try:
        from arm_sim_gazebo.sim_cameras import WORLD_LOOK_DOWN_COS, world_look_down_cos
        from arm_sim_gazebo import sim_scene

        scene = sim_scene.load_scene()
        b = scene.table_bounds
        tf = node.buf.lookup_transform("world", frame, rclpy.time.Time())
        x = tf.transform.translation.x
        y = tf.transform.translation.y
        margin = 0.05
        inside = (b.x_min - margin <= x <= b.x_max + margin) and (
            b.y_min - margin <= y <= b.y_max + margin
        )
        downward = world_look_down_cos(_tf_rot(node, "world", frame)) <= WORLD_LOOK_DOWN_COS
        return inside and downward
    except Exception:  # noqa: BLE001 —— 拿不到就不拦人
        return None


def expected_intrinsics(ns: str):
    """从 cameras.yaml 算出该相机的应有内参。拿不到就返回 None（不阻断检查）。"""
    try:
        from arm_sim_gazebo.sim_cameras import intrinsics, load_cameras
    except ImportError:
        return None
    try:
        for s in load_cameras():
            if s.ros_ns == ns:
                return s, intrinsics(s, s.width or 640, s.height or 480)
    except Exception:
        return None
    return None


def check_one(ns: str, args) -> tuple[bool, list[str]]:
    """检查一台相机。返回 (是否通过, 问题列表)。"""
    problems: list[str] = []
    node = Grab(ns)

    if not spin_until(
        node, lambda: node.img is not None and node.info is not None, args.timeout
    ):
        node.destroy_node()
        return False, [f"拿不到图像或内参（等 {args.timeout:.0f}s）"]

    frame = node.info.header.frame_id
    if not spin_until(
        node, lambda: node.buf.can_transform("world", frame, rclpy.time.Time()), 20.0
    ):
        node.destroy_node()
        return False, [f"TF world ← {frame} 查不到（相机在世上没有位姿）"]

    K = np.array(node.info.k, dtype=np.float64).reshape(3, 3)
    fx, fy, cx, cy = K[0, 0], K[1, 1], K[0, 2], K[1, 2]
    if not (frame.endswith("_optical_frame") or "optical" in frame):
        problems.append(
            f"frame_id 是 {frame}，看着不像光学系（应以 _optical_frame 结尾）。"
            f"盖成本体系时 tf2 查得到、不报错，但反投影出的点会整体转错。"
        )

    mm = np.frombuffer(node.img.data, dtype=np.uint16)
    if mm.size != node.img.height * node.img.width:
        node.destroy_node()
        return False, [
            f"深度图像不是 16UC1 单通道（{node.img.encoding}, "
            f"{len(node.img.data)} 字节 / 期望 {node.img.height * node.img.width * 2}）"
            f"—— gz 的 32FC1 米必须转成 16UC1 毫米，这一步没生效"
        ]
    mm = mm.reshape(node.img.height, node.img.width)

    # 与 cameras.yaml 推导出的应有内参比对：这能直接抓住"gz 的默认内参被透传"
    # 那个坑（它会给出 fx=277, cx=160, cy=120）。
    exp = expected_intrinsics(ns)
    if exp is not None:
        spec, (exp_k, _) = exp
        if abs(exp_k[0] - fx) > 1.0 or abs(exp_k[2] - cx) > 1.0:
            problems.append(
                f"内参不是按 cameras.yaml 算出来的：实测 fx={fx:.1f} cx={cx:.1f}，"
                f"应为 fx={exp_k[0]:.1f} cx={exp_k[2]:.1f}"
                f"（视场 {spec.horizontal_fov}）。\n"
                f"      这正是 gz 透传默认内参（fx=277, cx=160）时的症状，"
                f"会让所有三维点位偏半米量级。"
            )

    valid = (mm > 200) & (mm < 8000)
    vs, us = np.nonzero(valid)
    if len(us) < 500:
        node.destroy_node()
        return False, [f"有效深度像素太少（{len(us)} 个），图像基本是空的"]
    rng = np.random.default_rng(0)
    if len(us) > args.max_pixels:
        sel = rng.choice(len(us), args.max_pixels, replace=False)
        us, vs = us[sel], vs[sel]

    Z = mm[vs, us].astype(np.float64) / 1000.0
    pts = np.stack([(us - cx) * Z / fx, (vs - cy) * Z / fy, Z], axis=1)

    tf = node.buf.lookup_transform("world", frame, rclpy.time.Time())
    R = quat_to_R(
        (
            tf.transform.rotation.x,
            tf.transform.rotation.y,
            tf.transform.rotation.z,
            tf.transform.rotation.w,
        )
    )
    t = np.array(
        [
            tf.transform.translation.x,
            tf.transform.translation.y,
            tf.transform.translation.z,
        ]
    )
    zw = (pts @ R.T + t)[:, 2]
    # 安装朝向要在 destroy_node 之前查 —— 它还要用 TF（腕相机的视轴相对法兰）
    mount_problems(node, ns, frame, problems)
    above_table = looking_at_table(node, frame)
    node.destroy_node()

    print(f"\n[{ns}]  光学系={frame}   fx={fx:.1f} fy={fy:.1f} cx={cx:.1f} cy={cy:.1f}")
    print(f"        world-z 中位={np.median(zw):+.3f} m  范围[{zw.min():+.2f}, {zw.max():+.2f}]")

    # 逐平面命中率
    rates = {}
    for pz, name in args.planes:
        rate = float(np.mean(np.abs(zw - pz) < args.tol)) * 100.0
        rates[pz] = rate
        print(f"        命中 {name}(z={pz:+.2f}) : {rate:5.1f}%")

    # 期望判据
    table_z = 0.40
    for want_z, want_min in args.expect.get(ns, []):
        name = dict(args.planes).get(want_z, f"z={want_z}")
        got = rates.get(want_z, 0.0)
        if got >= want_min:
            continue
        if abs(want_z - table_z) < 1e-9 and above_table is False:
            # 腕相机不在台面上方时，"没看见台面"是正常的（它看见的是地面）。
            print(
                f"        ⚠ 跳过「{name}」判据：相机此刻没朝下看着台面"
                f"（光心不在 scene.yaml 的 table_bounds 内，或视轴没朝下）——\n"
                f"          此刻只看得见地面是正常的。\n"
                f"          想让这条生效：把机械臂摆到工件上方（例如 S04 跟踪时）再跑本脚本。"
            )
            continue
        problems.append(
            f"只看见 {name} 的 {got:.1f}%，要求 ≥{want_min:.0f}%。\n"
            f"      相机很可能是**背对**目标（位姿 rpy 写错）—— "
            f"这种错几何自洽、不报错，只会让相机拍不到该拍的东西。"
        )

    total = sum(rates.values())
    if total < 1.0:
        problems.append(f"所有预期平面加起来命中率只有 {total:.1f}%，几乎没有指向场景")

    ok = not problems
    print(f"       → {'✓ 通过' if ok else '✘ 不通过'}")
    for p in problems:
        print(f"       ✘ {p}")
    return ok, problems


def parse_planes(text: str):
    out = []
    for item in text.split(","):
        item = item.strip()
        if not item:
            continue
        if ":" in item:
            z, name = item.split(":", 1)
        else:
            z, name = item, f"z={item}"
        out.append((float(z), name.strip()))
    if not out:
        raise argparse.ArgumentTypeError("--planes 不能为空")
    return out


def parse_expect(items):
    """把 'd435:0.40:10' 解析成 {ns: [(z, min_pct), ...]}。"""
    out: dict[str, list[tuple[float, float]]] = {}
    for it in items:
        parts = it.split(":")
        if len(parts) < 2:
            raise argparse.ArgumentTypeError(f"--expect 格式应为 相机:z[:最低百分比]，收到 {it!r}")
        ns = parts[0]
        z = float(parts[1])
        pct = float(parts[2]) if len(parts) > 2 else 5.0
        out.setdefault(ns, []).append((z, pct))
    return out


def main() -> int:
    ap = argparse.ArgumentParser(
        description="仿真相机几何验收：反投影深度，检查是否命中预期平面。",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    ap.add_argument("--cameras", default="sim_d405,sim_d435", help="要检查的相机命名空间，逗号分隔")
    ap.add_argument(
        "--planes",
        type=parse_planes,
        default=parse_planes("0.0:地面,0.40:台面"),
        help="要统计命中率的平面，格式 z:名称，逗号分隔（默认 0.0:地面,0.40:台面）",
    )
    ap.add_argument(
        "--expect",
        action="append",
        dest="expect_raw",
        default=None,
        help="要求某相机命中某平面，格式 相机:z[:最低百分比]，可重复。"
        f"默认 {list(DEFAULT_EXPECT)}",
    )
    ap.add_argument("--tol", type=float, default=0.02, help="平面命中容差（米），默认 0.02")
    ap.add_argument("--timeout", type=float, default=25.0, help="等图像/内参的秒数")
    ap.add_argument("--max-pixels", type=int, default=40000, help="最多采样多少像素")
    args = ap.parse_args()
    args.expect = parse_expect(args.expect_raw or list(DEFAULT_EXPECT))

    rclpy.init()
    try:
        results = {}
        for ns in [c.strip() for c in args.cameras.split(",") if c.strip()]:
            try:
                results[ns] = check_one(ns, args)
            except KeyboardInterrupt:
                raise
            except Exception as exc:  # 单台相机出错不该拖垮整轮检查
                print(f"\n[{ns}] ✘ 检查时异常：{type(exc).__name__}: {exc}")
                results[ns] = (False, [str(exc)])
    finally:
        rclpy.try_shutdown()

    failed = [ns for ns, (ok, _) in results.items() if not ok]
    print()
    if failed:
        print(f"=== ✘ 不通过：{', '.join(failed)} ===")
        print("    仿真里相机出错几乎都不报错，别只看'有没有图像'。")
        return 1
    print(f"=== ✓ 全部通过：{', '.join(results)} ===")
    return 0


if __name__ == "__main__":
    sys.exit(main())