#!/usr/bin/env python3
"""AprilTag map localization for Apollo tag navigation mode.

This node consumes the vision pipeline's JSON tag detections and a saved
``*.tags.yaml`` map. It publishes ``map -> odom`` from the best visible mapped
tag while wheel odometry and optional IMU yaw continue to handle motion between
tag corrections.
"""

import json
import math
import os
import re
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple

import rclpy
from geometry_msgs.msg import PoseWithCovarianceStamped, TransformStamped
from rclpy.node import Node
from sensor_msgs.msg import Imu
from std_msgs.msg import String
import tf2_ros

try:
    import yaml
except Exception:
    yaml = None


@dataclass(frozen=True)
class Pose2D:
    x: float
    y: float
    yaw: float


@dataclass(frozen=True)
class Transform2D:
    """Planar rigid transform: parent_T_child."""

    x: float
    y: float
    yaw: float

    def compose(self, child: "Transform2D") -> "Transform2D":
        cos_y = math.cos(self.yaw)
        sin_y = math.sin(self.yaw)
        return Transform2D(
            self.x + cos_y * child.x - sin_y * child.y,
            self.y + sin_y * child.x + cos_y * child.y,
            normalize_yaw(self.yaw + child.yaw),
        )

    def inverse(self) -> "Transform2D":
        cos_y = math.cos(self.yaw)
        sin_y = math.sin(self.yaw)
        return Transform2D(
            -cos_y * self.x - sin_y * self.y,
            sin_y * self.x - cos_y * self.y,
            normalize_yaw(-self.yaw),
        )

    def as_pose(self) -> Pose2D:
        return Pose2D(self.x, self.y, normalize_yaw(self.yaw))


@dataclass(frozen=True)
class TagObservation:
    tag_id: int
    forward: float
    right: float
    vertical: float
    yaw_rad: float
    pitch_deg: Optional[float]
    roll_deg: Optional[float]
    yaw_deg: float
    distance: float
    reprojection_error: Optional[float]
    camera_model: str
    pose_image_source: str


def yaw_to_quaternion(yaw: float) -> Tuple[float, float, float, float]:
    half = yaw * 0.5
    return 0.0, 0.0, math.sin(half), math.cos(half)


def quaternion_to_yaw(q) -> float:
    siny_cosp = 2.0 * (q.w * q.z + q.x * q.y)
    cosy_cosp = 1.0 - 2.0 * (q.y * q.y + q.z * q.z)
    return math.atan2(siny_cosp, cosy_cosp)


def normalize_yaw(yaw: float) -> float:
    return math.atan2(math.sin(yaw), math.cos(yaw))


def finite_float(value: object) -> Optional[float]:
    try:
        number = float(value)
    except (TypeError, ValueError):
        return None
    return number if math.isfinite(number) else None


def env_bool(name: str, default: bool) -> bool:
    value = os.environ.get(name)
    if value is None:
        return default
    return value.strip().lower() in {"1", "true", "yes", "on"}


def env_float(name: str, default: float) -> float:
    value = finite_float(os.environ.get(name))
    return default if value is None else value


def pose_to_transform(pose: Pose2D) -> Transform2D:
    return Transform2D(pose.x, pose.y, pose.yaw)


def camera_tag_from_observation(observation: TagObservation) -> Transform2D:
    return Transform2D(
        observation.forward,
        -observation.right,
        observation.yaw_rad,
    )


def robot_pose_from_tag(
    tag_pose: Pose2D,
    observation: TagObservation,
    base_from_camera: Transform2D,
) -> Pose2D:
    map_from_tag = pose_to_transform(tag_pose)
    camera_from_tag = camera_tag_from_observation(observation)
    base_from_tag = base_from_camera.compose(camera_from_tag)
    return map_from_tag.compose(base_from_tag.inverse()).as_pose()

def robot_yaw_from_tag_yaw(
    tag_pose: Pose2D,
    observation: TagObservation,
    base_from_camera: Transform2D,
) -> float:
    return normalize_yaw(
        tag_pose.yaw + observation.yaw_rad + math.pi - base_from_camera.yaw
    )
    
def robot_pose_from_tag_direct(
    tag_pose: Pose2D,
    observation: TagObservation,
    base_from_camera: Transform2D,
) -> Pose2D:
    robot_yaw = robot_yaw_from_tag_yaw(tag_pose, observation, base_from_camera)

    # tag position relative to camera, using your pipeline convention:
    # +x forward, +right right
    cam_tag_x = observation.forward
    cam_tag_y = -observation.right

    # include camera position on robot
    base_tag_x = base_from_camera.x + cam_tag_x
    base_tag_y = base_from_camera.y + cam_tag_y

    c = math.cos(robot_yaw)
    s = math.sin(robot_yaw)

    map_dx = c * base_tag_x - s * base_tag_y
    map_dy = s * base_tag_x + c * base_tag_y

    return Pose2D(
        tag_pose.x - map_dx,
        tag_pose.y - map_dy,
        robot_yaw,
    )

def map_odom_from_robot_pose(robot_map: Pose2D, odom_base: Pose2D) -> Pose2D:
    return pose_to_transform(robot_map).compose(pose_to_transform(odom_base).inverse()).as_pose()


def robot_pose_from_map_odom(map_odom: Pose2D, odom_base: Pose2D) -> Pose2D:
    return pose_to_transform(map_odom).compose(pose_to_transform(odom_base)).as_pose()


def read_scalar_yaml(path: str) -> Dict[str, object]:
    data: Dict[str, object] = {}
    try:
        with open(path, "r", encoding="utf-8") as f:
            for line in f:
                stripped = line.strip()
                if not stripped or stripped.startswith("#") or ":" not in stripped:
                    continue
                key, value = stripped.split(":", 1)
                value = value.strip()
                if key == "resolution":
                    data[key] = float(value)
                elif key == "origin":
                    nums = re.findall(r"[-+]?(?:\d*\.\d+|\d+)(?:[eE][-+]?\d+)?", value)
                    data[key] = [float(v) for v in nums[:3]]
    except Exception:
        pass
    return data


def pgm_height_from_yaml(map_yaml: str) -> Optional[int]:
    try:
        image_rel = None
        with open(map_yaml, "r", encoding="utf-8") as f:
            for line in f:
                stripped = line.strip()
                if stripped.startswith("image:"):
                    image_rel = stripped.split(":", 1)[1].strip().strip('"')
                    break
        if not image_rel:
            return None
        pgm_path = image_rel if os.path.isabs(image_rel) else os.path.join(os.path.dirname(map_yaml), image_rel)
        with open(pgm_path, "rb") as f:
            tokens = []
            while len(tokens) < 3:
                line = f.readline()
                if not line:
                    break
                line = line.split(b"#", 1)[0]
                tokens.extend(line.split())
        if len(tokens) >= 3:
            return int(tokens[2])
    except Exception:
        return None
    return None


def read_tag_map(path: str, map_yaml: str) -> Dict[int, Pose2D]:
    map_info = read_scalar_yaml(map_yaml)
    resolution = float(map_info.get("resolution", 0.05))
    origin = map_info.get("origin", [0.0, 0.0, 0.0])
    if not isinstance(origin, list) or len(origin) < 2:
        origin = [0.0, 0.0, 0.0]
    image_height = pgm_height_from_yaml(map_yaml)

    def add_tag(tags: Dict[int, Pose2D], raw: Dict[str, object]) -> None:
        if not isinstance(raw, dict) or "id" not in raw or "x" not in raw or "y" not in raw:
            return
        tag_id = int(raw["id"])
        px = float(raw["x"])
        py = float(raw["y"])
        yaw = math.radians(float(raw.get("theta", raw.get("yaw", 0.0))))
        world_x = float(origin[0]) + px * resolution
        if image_height is not None:
            world_y = float(origin[1]) + (image_height - 1 - py) * resolution
        else:
            world_y = float(origin[1]) + py * resolution
        tags[tag_id] = Pose2D(world_x, world_y, yaw)

    tags: Dict[int, Pose2D] = {}
    if yaml is not None:
        try:
            with open(path, "r", encoding="utf-8") as f:
                data = yaml.safe_load(f) or {}
            raw_tags = data.get("tags", []) if isinstance(data, dict) else data
            if isinstance(raw_tags, list):
                for raw in raw_tags:
                    add_tag(tags, raw)
                return tags
        except Exception:
            tags.clear()

    current: Dict[str, float] = {}

    def flush() -> None:
        add_tag(tags, current)

    try:
        with open(path, "r", encoding="utf-8") as f:
            for raw in f:
                line = raw.strip()
                if not line:
                    continue
                if line.startswith("- "):
                    flush()
                    current = {}
                    line = line[2:].strip()
                if ":" in line:
                    key, value = line.split(":", 1)
                    key = key.strip()
                    value = value.strip().strip('"')
                    if key in {"id", "x", "y", "theta", "yaw"}:
                        try:
                            current[key] = float(value)
                        except ValueError:
                            pass
        flush()
    except Exception:
        return {}
    return tags


def read_tag_config(path: str) -> Tuple[Optional[float], Optional[str]]:
    if yaml is not None:
        try:
            with open(path, "r", encoding="utf-8") as f:
                data = yaml.safe_load(f) or {}
            if isinstance(data, dict):
                size = data.get("tag_size")
                family = data.get("tag_family")
                size_f = float(size) if size is not None and float(size) > 0 else None
                family_s = str(family).strip() if family else None
                return size_f, family_s
        except Exception:
            pass

    size_val: Optional[float] = None
    family_val: Optional[str] = None
    try:
        with open(path, "r", encoding="utf-8") as f:
            for line in f:
                stripped = line.strip()
                if not stripped or stripped.startswith("#") or ":" not in stripped:
                    continue
                if stripped.startswith("- "):
                    break
                key, _, value = stripped.partition(":")
                key = key.strip()
                value = value.strip().strip('"')
                if key == "tag_size":
                    try:
                        v = float(value)
                        if v > 0:
                            size_val = v
                    except ValueError:
                        pass
                elif key == "tag_family" and value:
                    family_val = value
    except Exception:
        pass
    return size_val, family_val


class TagLocalizationNode(Node):
    def __init__(self) -> None:
        super().__init__("tag_localization_node")

        self.map_yaml = os.environ.get("APOLLO_MAP_YAML", "")
        self.tags_yaml = os.environ.get("APOLLO_TAGS_YAML", "")
        self.map_frame = str(self.declare_parameter("map_frame", os.environ.get("APOLLO_TAG_MAP_FRAME", "map")).value)
        self.odom_frame = str(self.declare_parameter("odom_frame", os.environ.get("APOLLO_TAG_ODOM_FRAME", "odom")).value)
        self.base_frame = str(
            self.declare_parameter("base_frame", os.environ.get("APOLLO_TAG_BASE_FRAME", "base_footprint")).value
        )
        self.use_odom = bool(self.declare_parameter("use_odom", env_bool("APOLLO_TAG_USE_ODOM", True)).value)
        self.use_imu = bool(self.declare_parameter("use_imu", env_bool("APOLLO_TAG_USE_IMU", True)).value)
        self.use_imu_yaw = self.use_imu
        self.imu_topic = str(self.declare_parameter("imu_topic", os.environ.get("APOLLO_TAG_IMU_TOPIC", "imu")).value)
        self.imu_timeout = float(
            self.declare_parameter("imu_timeout", env_float("APOLLO_TAG_IMU_TIMEOUT", 0.75)).value
        )

        self.max_abs_yaw_deg = float(
            self.declare_parameter("max_abs_yaw_deg", env_float("APOLLO_TAG_MAX_ABS_YAW_DEG", 55.0)).value
        )
        self.max_abs_pitch_deg = float(
            self.declare_parameter("max_abs_pitch_deg", env_float("APOLLO_TAG_MAX_ABS_PITCH_DEG", 45.0)).value
        )
        self.max_abs_roll_deg = float(
            self.declare_parameter("max_abs_roll_deg", env_float("APOLLO_TAG_MAX_ABS_ROLL_DEG", 45.0)).value
        )
        self.max_reprojection_error = float(
            self.declare_parameter(
                "max_reprojection_error",
                env_float("APOLLO_TAG_MAX_REPROJECTION_ERROR", 5.0),
            ).value
        )
        self.max_correction_distance = float(
            self.declare_parameter(
                "max_correction_distance",
                env_float("APOLLO_TAG_MAX_CORRECTION_DISTANCE", 0.75),
            ).value
        )
        self.max_correction_yaw_deg = float(
            self.declare_parameter(
                "max_correction_yaw_deg",
                env_float("APOLLO_TAG_MAX_CORRECTION_YAW_DEG", 45.0),
            ).value
        )
        self.correction_alpha_pos = float(
            self.declare_parameter(
                "correction_alpha_pos",
                env_float("APOLLO_TAG_CORRECTION_ALPHA_POS", 0.35),
            ).value
        )
        if not 0.0 < self.correction_alpha_pos <= 1.0:
            self.get_logger().warn("correction_alpha_pos must be in (0, 1]; using 0.35")
            self.correction_alpha_pos = 0.35
        self.correction_alpha_yaw = float(
            self.declare_parameter(
                "correction_alpha_yaw",
                env_float("APOLLO_TAG_CORRECTION_ALPHA_YAW", 0.35),
            ).value
        )
        if not 0.0 < self.correction_alpha_yaw <= 1.0:
            self.get_logger().warn("correction_alpha_yaw must be in (0, 1]; using 0.35")
            self.correction_alpha_yaw = 0.35
        self.min_stable_frames = max(
            1,
            int(
                self.declare_parameter(
                    "min_stable_frames",
                    int(os.environ.get("APOLLO_TAG_MIN_STABLE_FRAMES", "2")),
                ).value
            ),
        )

        camera_x = float(self.declare_parameter("camera_x", env_float("APOLLO_TAG_CAMERA_X", 0.0)).value)
        camera_y = float(self.declare_parameter("camera_y", env_float("APOLLO_TAG_CAMERA_Y", 0.0)).value)
        camera_yaw_deg = float(
            self.declare_parameter("camera_yaw_deg", env_float("APOLLO_TAG_CAMERA_YAW_DEG", 0.0)).value
        )
        self.base_from_camera = Transform2D(camera_x, camera_y, math.radians(camera_yaw_deg))
        self.vision_yaw_sign = float(
            self.declare_parameter("vision_yaw_sign", env_float("APOLLO_TAG_VISION_YAW_SIGN", 1.0)).value
        )
        vision_yaw_offset_deg = float(
            self.declare_parameter(
                "vision_yaw_offset_deg",
                env_float("APOLLO_TAG_VISION_YAW_OFFSET_DEG", 0.0),
            ).value
        )
        self.vision_yaw_offset = math.radians(vision_yaw_offset_deg)
        self.pnp_forward_scale = float(
            self.declare_parameter(
                "pnp_forward_scale",
                env_float("APOLLO_TAG_PNP_FORWARD_SCALE", 1.0),
            ).value
        )
        if self.pnp_forward_scale <= 0.0:
            self.get_logger().warn("pnp_forward_scale must be > 0; using 1.0")
            self.pnp_forward_scale = 1.0
        self.pnp_sideways_scale = float(
            self.declare_parameter(
                "pnp_sideways_scale",
                env_float("APOLLO_TAG_PNP_SIDEWAYS_SCALE", 1.0),
            ).value
        )
        if self.pnp_sideways_scale <= 0.0:
            self.get_logger().warn("pnp_sideways_scale must be > 0; using 1.0")
            self.pnp_sideways_scale = 1.0

        self.tag_map = read_tag_map(self.tags_yaml, self.map_yaml) if self.tags_yaml and self.map_yaml else {}
        self._expected_tag_size: Optional[float] = None
        if self.tags_yaml:
            self._expected_tag_size, _ = read_tag_config(self.tags_yaml)

        self.tf_buffer = tf2_ros.Buffer()
        self.tf_listener = tf2_ros.TransformListener(self.tf_buffer, self)
        self.tf_broadcaster = tf2_ros.TransformBroadcaster(self)
        self.source_pub = self.create_publisher(String, "/apollo/tag_localization/source", 10)
        self.sub = self.create_subscription(String, "/apollo/vision/output", self.on_vision, 10)
        self.initial_pose_sub = self.create_subscription(
            PoseWithCovarianceStamped,
            "/initialpose",
            self.on_initial_pose,
            10,
        )
        self.imu_sub = None
        if self.use_imu and self.imu_topic:
            self.imu_sub = self.create_subscription(Imu, self.imu_topic, self.on_imu, 20)
        self.timer = self.create_timer(0.05, self.on_timer)

        self.map_odom = Pose2D(0.0, 0.0, 0.0)
        self.localization_source = "odom_imu_fallback"
        self.pending_initial_pose: Optional[PoseWithCovarianceStamped] = None
        self.latest_imu_yaw: Optional[float] = None
        self.latest_imu_time = None
        self.imu_yaw_offset: Optional[float] = None
        self.warned_imu_wait = False
        self.warned_empty_tag_map = False
        self.warned_unknown_tags = False
        self.warned_initial_pose_wait = False
        self.warned_tag_size_mismatch = False
        self.last_best_tag_id: Optional[int] = None
        self.stable_count = 0
        self.pose_initialized = False
        # True only after a real AprilTag fix has been accepted. The anti-jump
        # gates apply between tag fixes; the first tag fix is allowed to fully
        # override a rough /initialpose prior regardless of correction size.
        self.tag_fix_received = False

        self.log_startup()
        if not self.tag_map:
            self.get_logger().warn(
                "tag_localization_node loaded 0 tags; map->odom will remain odom/IMU-only "
                "until the active map has saved tags with id/x/y entries."
            )

    def log_startup(self) -> None:
        self.get_logger().info(
            "tag_localization_node ready: "
            f"tags={len(self.tag_map)} map={self.map_yaml} tags_yaml={self.tags_yaml} "
            f"map_frame={self.map_frame} odom_frame={self.odom_frame} base_frame={self.base_frame} "
            f"use_odom={self.use_odom} use_imu={self.use_imu} imu_topic={self.imu_topic} "
            f"imu_timeout={self.imu_timeout:.2f}"
        )
        self.get_logger().info(
            "AprilTag localization params: "
            "translation=pnp "
            f"{self.localization_gate_summary()} "
            "(0.0 disables a gate)"
        )
        self.get_logger().info(
            "AprilTag frame params: "
            f"camera_offset=({self.base_from_camera.x:.3f},{self.base_from_camera.y:.3f},"
            f"{math.degrees(self.base_from_camera.yaw):.1f}deg) "
            f"vision_yaw_sign={self.vision_yaw_sign:.2f} "
            f"vision_yaw_offset_deg={math.degrees(self.vision_yaw_offset):.1f} "
            f"pnp_forward_scale={self.pnp_forward_scale:.3f} "
            f"pnp_sideways_scale={self.pnp_sideways_scale:.3f}"
        )

    def localization_gate_summary(self) -> str:
        return (
            f"max_abs_yaw_deg={self.max_abs_yaw_deg:.1f} "
            f"max_abs_pitch_deg={self.max_abs_pitch_deg:.1f} "
            f"max_abs_roll_deg={self.max_abs_roll_deg:.1f} "
            f"max_reprojection_error={self.max_reprojection_error:.2f} "
            f"max_correction_distance={self.max_correction_distance:.2f} "
            f"max_correction_yaw_deg={self.max_correction_yaw_deg:.1f} "
            f"correction_alpha_pos={self.correction_alpha_pos:.2f} "
            f"correction_alpha_yaw={self.correction_alpha_yaw:.2f} "
            f"min_stable_frames={self.min_stable_frames}"
        )

    def reject_observation(
        self,
        tag_id: int,
        reason: str,
        yaw_deg: Optional[float],
        pitch_deg: Optional[float],
        roll_deg: Optional[float],
        distance: Optional[float],
        reprojection_error: Optional[float],
    ) -> None:
        self.get_logger().debug(
            f"Tag {tag_id}: rejected - {reason}; "
            f"yaw={yaw_deg} pitch={pitch_deg} roll={roll_deg} "
            f"distance={distance} reprojectionError={reprojection_error}; "
            f"{self.localization_gate_summary()}"
        )

    def observation_from_detection(self, detection: Dict[str, object]) -> Optional[TagObservation]:
        try:
            tag_id = int(detection.get("id", -1))
        except (TypeError, ValueError):
            return None

        forward = finite_float(detection.get("pnpForward"))
        right = finite_float(detection.get("pnpRight"))
        vertical = finite_float(detection.get("pnpUp"))
        yaw_deg = finite_float(detection.get("yaw"))
        pitch_deg = finite_float(detection.get("pitch"))
        roll_deg = finite_float(detection.get("roll"))
        reprojection_error = finite_float(detection.get("reprojectionError"))

        distance = None
        if forward is not None and right is not None and vertical is not None:
            distance = math.sqrt(forward * forward + right * right + vertical * vertical)

        if tag_id not in self.tag_map:
            self.reject_observation(tag_id, "tag id is not in map", yaw_deg, pitch_deg, roll_deg, distance, reprojection_error)
            return None
        if forward is None or forward <= 0.0:
            self.reject_observation(tag_id, "pnpForward missing or <= 0", yaw_deg, pitch_deg, roll_deg, distance, reprojection_error)
            return None
        if right is None or vertical is None:
            self.reject_observation(tag_id, "pnpRight or pnpUp missing", yaw_deg, pitch_deg, roll_deg, distance, reprojection_error)
            return None
        if yaw_deg is None:
            self.reject_observation(tag_id, "yaw missing", yaw_deg, pitch_deg, roll_deg, distance, reprojection_error)
            return None
        if self.max_abs_yaw_deg > 0.0 and abs(yaw_deg) > self.max_abs_yaw_deg:
            self.reject_observation(tag_id, "yaw over gate", yaw_deg, pitch_deg, roll_deg, distance, reprojection_error)
            return None
        if self.max_abs_pitch_deg > 0.0 and pitch_deg is not None and abs(pitch_deg) > self.max_abs_pitch_deg:
            self.reject_observation(tag_id, "pitch over gate", yaw_deg, pitch_deg, roll_deg, distance, reprojection_error)
            return None
        if self.max_abs_roll_deg > 0.0 and roll_deg is not None and abs(roll_deg) > self.max_abs_roll_deg:
            self.reject_observation(tag_id, "roll over gate", yaw_deg, pitch_deg, roll_deg, distance, reprojection_error)
            return None
        if (
            self.max_reprojection_error > 0.0
            and reprojection_error is not None
            and reprojection_error > self.max_reprojection_error
        ):
            self.reject_observation(
                tag_id,
                "reprojectionError over gate",
                yaw_deg,
                pitch_deg,
                roll_deg,
                distance,
                reprojection_error,
            )
            return None

        yaw_rad = normalize_yaw(self.vision_yaw_sign * math.radians(yaw_deg) + self.vision_yaw_offset)
        scaled_forward = forward * self.pnp_forward_scale
        scaled_right = right * self.pnp_sideways_scale
        return TagObservation(
            tag_id=tag_id,
            forward=scaled_forward,
            right=scaled_right,
            vertical=vertical,
            yaw_rad=yaw_rad,
            pitch_deg=pitch_deg,
            roll_deg=roll_deg,
            yaw_deg=yaw_deg,
            distance=math.sqrt(scaled_forward * scaled_forward + scaled_right * scaled_right + vertical * vertical),
            reprojection_error=reprojection_error,
            camera_model=str(detection.get("cameraModel", "")),
            pose_image_source=str(detection.get("poseImageSource", "")),
        )

    def best_observation(self, detections: object) -> Tuple[Optional[TagObservation], List[int]]:
        if not isinstance(detections, list):
            return None, []

        observations: List[TagObservation] = []
        unknown_ids: List[int] = []
        for detection in detections:
            if not isinstance(detection, dict):
                continue
            try:
                tag_id = int(detection.get("id", -1))
            except (TypeError, ValueError):
                continue
            if tag_id not in self.tag_map:
                unknown_ids.append(tag_id)
                self.observation_from_detection(detection)
                continue
            observation = self.observation_from_detection(detection)
            if observation is not None:
                observations.append(observation)

        if not observations:
            return None, unknown_ids

        observations.sort(
            key=lambda obs: (
                abs(obs.yaw_deg) / 45.0 + obs.distance / 2.0,
                obs.reprojection_error if obs.reprojection_error is not None else float("inf"),
            )
        )
        return observations[0], unknown_ids

    def warn_on_tag_size_mismatch(self, detections: object) -> None:
        if self.warned_tag_size_mismatch or self._expected_tag_size is None or not isinstance(detections, list):
            return
        for detection in detections:
            if not isinstance(detection, dict):
                continue
            observed_size = finite_float(detection.get("tagSizeMeters"))
            if observed_size is None:
                continue
            if abs(observed_size - self._expected_tag_size) > 0.001:
                self.get_logger().warn(
                    f"Vision is solving AprilTags with tagSizeMeters={observed_size:.4f}, "
                    f"but the loaded tag map expects tag_size={self._expected_tag_size:.4f}. "
                    "Fix the tag map/vision config mismatch before trusting distance."
                )
                self.warned_tag_size_mismatch = True
                return

    def on_imu(self, msg: Imu) -> None:
        q = msg.orientation
        norm = math.sqrt(q.x * q.x + q.y * q.y + q.z * q.z + q.w * q.w)
        if norm < 1e-6:
            return
        self.latest_imu_yaw = quaternion_to_yaw(q)
        self.latest_imu_time = self.get_clock().now()
        self.warned_imu_wait = False

    def fresh_imu_yaw(self) -> Optional[float]:
        if not self.use_imu_yaw or self.latest_imu_yaw is None or self.latest_imu_time is None:
            if self.use_imu_yaw and not self.warned_imu_wait:
                self.get_logger().warn(
                    f"Waiting for IMU yaw on '{self.imu_topic}'; using wheel odom heading until it arrives."
                )
                self.warned_imu_wait = True
            return None

        age = (self.get_clock().now() - self.latest_imu_time).nanoseconds / 1e9
        if age > self.imu_timeout:
            if not self.warned_imu_wait:
                self.get_logger().warn(
                    f"IMU yaw on '{self.imu_topic}' is stale ({age:.2f}s); using wheel odom heading."
                )
                self.warned_imu_wait = True
            return None
        return self.latest_imu_yaw

    def update_imu_reference(self, robot_map_yaw: float) -> None:
        imu_yaw = self.fresh_imu_yaw()
        if imu_yaw is None:
            return
        first_reference = self.imu_yaw_offset is None
        self.imu_yaw_offset = normalize_yaw(robot_map_yaw - imu_yaw)
        if first_reference:
            self.get_logger().info("IMU yaw reference initialized from tag/map localization.")

    def lookup_odom_base(self) -> Optional[TransformStamped]:
        try:
            return self.tf_buffer.lookup_transform(
                self.odom_frame,
                self.base_frame,
                rclpy.time.Time(),
                timeout=rclpy.duration.Duration(seconds=0.05),
            )
        except Exception as exc:
            self.get_logger().debug(f"Waiting for odom -> base transform: {exc}")
            return None

    def pose_from_transform(self, tf: TransformStamped) -> Pose2D:
        return Pose2D(
            float(tf.transform.translation.x),
            float(tf.transform.translation.y),
            quaternion_to_yaw(tf.transform.rotation),
        )

    def set_map_odom(
        self,
        pose: Pose2D,
        source: str,
        robot_map_yaw: Optional[float] = None,
    ) -> None:
        previous_source = self.localization_source
        self.map_odom = Pose2D(pose.x, pose.y, normalize_yaw(pose.yaw))
        self.localization_source = source
        if source == "initialpose" or source.startswith("tag:"):
            self.pose_initialized = True
        if source.startswith("tag:"):
            self.tag_fix_received = True
        if robot_map_yaw is not None:
            self.update_imu_reference(robot_map_yaw)
        if source.startswith("tag:"):
            self.pending_initial_pose = None
            self.warned_initial_pose_wait = False
        if source != previous_source:
            self.get_logger().info(f"Localization source: {source}")
        self.broadcast_map_odom()

    def broadcast_map_odom(self) -> None:
        tf = TransformStamped()
        tf.header.stamp = self.get_clock().now().to_msg()
        tf.header.frame_id = self.map_frame
        tf.child_frame_id = self.odom_frame
        tf.transform.translation.x = self.map_odom.x
        tf.transform.translation.y = self.map_odom.y
        tf.transform.translation.z = 0.0
        qx, qy, qz, qw = yaw_to_quaternion(self.map_odom.yaw)
        tf.transform.rotation.x = qx
        tf.transform.rotation.y = qy
        tf.transform.rotation.z = qz
        tf.transform.rotation.w = qw
        self.tf_broadcaster.sendTransform(tf)

    def apply_imu_yaw_to_map_odom(self) -> None:
        if self.imu_yaw_offset is None:
            return
        imu_yaw = self.fresh_imu_yaw()
        if imu_yaw is None:
            return
        odom_base = self.lookup_odom_base()
        if odom_base is None:
            return

        odom_base_pose = self.pose_from_transform(odom_base)
        robot_map_pose = robot_pose_from_map_odom(self.map_odom, odom_base_pose)
        imu_robot_pose = Pose2D(
            robot_map_pose.x,
            robot_map_pose.y,
            normalize_yaw(imu_yaw + self.imu_yaw_offset),
        )
        self.map_odom = map_odom_from_robot_pose(imu_robot_pose, odom_base_pose)

    def on_timer(self) -> None:
        if self.pending_initial_pose is not None:
            self.apply_initial_pose(self.pending_initial_pose, warn_if_unavailable=False)
        if self.use_odom:
            self.apply_imu_yaw_to_map_odom()
        self.broadcast_map_odom()
        self.source_pub.publish(String(data=self.localization_source))

    def on_initial_pose(self, msg: PoseWithCovarianceStamped) -> None:
        if self.localization_source.startswith("tag:"):
            self.get_logger().info(
                f"Ignoring /initialpose because tag localization is already initialized ({self.localization_source})."
            )
            return

        frame_id = (msg.header.frame_id or self.map_frame).strip()
        if frame_id != self.map_frame:
            self.get_logger().warn(
                f"Ignoring /initialpose in frame '{frame_id}'; expected '{self.map_frame}'."
            )
            return

        self.pending_initial_pose = msg
        self.apply_initial_pose(msg, warn_if_unavailable=True)

    def apply_initial_pose(
        self,
        msg: PoseWithCovarianceStamped,
        warn_if_unavailable: bool,
    ) -> bool:
        odom_base = self.lookup_odom_base()
        if odom_base is None:
            if warn_if_unavailable and not self.warned_initial_pose_wait:
                self.get_logger().warn("Cannot apply /initialpose yet; odom -> base transform is not available.")
                self.warned_initial_pose_wait = True
            return False

        pose = msg.pose.pose
        robot_map = Pose2D(
            pose.position.x,
            pose.position.y,
            quaternion_to_yaw(pose.orientation),
        )
        map_odom = map_odom_from_robot_pose(robot_map, self.pose_from_transform(odom_base))
        self.set_map_odom(map_odom, "initialpose", robot_map_yaw=robot_map.yaw)
        self.pending_initial_pose = None
        self.warned_initial_pose_wait = False
        self.last_best_tag_id = None
        self.stable_count = 0
        return True

    def on_vision(self, msg: String) -> None:
        try:
            payload = json.loads(msg.data)
            detections = payload.get("tags", [])
        except Exception:
            return

        self.warn_on_tag_size_mismatch(detections)

        if detections and not self.tag_map and not self.warned_empty_tag_map:
            ids = [int(d.get("id", -1)) for d in detections if isinstance(d, dict)]
            self.get_logger().warn(
                f"AprilTags detected {ids}, but the loaded map tag list is empty; "
                "save tags on the active map so map->odom can be published."
            )
            self.warned_empty_tag_map = True

        observation, unknown_ids = self.best_observation(detections)
        if observation is None:
            self.last_best_tag_id = None
            self.stable_count = 0
            if unknown_ids and self.tag_map and not self.warned_unknown_tags:
                self.get_logger().warn(
                    f"AprilTags detected {sorted(set(unknown_ids))}, but none are in the loaded map tag IDs "
                    f"{sorted(self.tag_map.keys())}."
                )
                self.warned_unknown_tags = True
            return

        if observation.tag_id != self.last_best_tag_id:
            self.last_best_tag_id = observation.tag_id
            self.stable_count = 1
        else:
            self.stable_count += 1
        if self.stable_count < self.min_stable_frames:
            self.get_logger().debug(
                f"Tag {observation.tag_id}: waiting for stability "
                f"({self.stable_count}/{self.min_stable_frames})"
            )
            return

        odom_base = self.lookup_odom_base()
        if odom_base is None:
            return

        odom_base_pose = self.pose_from_transform(odom_base)
        tag_pose = self.tag_map[observation.tag_id]
        #candidate = robot_pose_from_tag(tag_pose, observation, self.base_from_camera)
        candidate = robot_pose_from_tag_direct(
            tag_pose,
            observation,
            self.base_from_camera,
        )
        current = robot_pose_from_map_odom(self.map_odom, odom_base_pose)

        dx = candidate.x - current.x
        dy = candidate.y - current.y
        correction_distance = math.sqrt(dx * dx + dy * dy)
        correction_yaw = abs(math.degrees(normalize_yaw(candidate.yaw - current.yaw)))

        if self.tag_fix_received:
            if self.max_correction_distance > 0.0 and correction_distance > self.max_correction_distance:
                self.get_logger().warn(
                    f"Tag {observation.tag_id}: jump rejected - "
                    f"translation {correction_distance:.2f} m > {self.max_correction_distance:.2f} m; "
                    f"{self.localization_gate_summary()}"
                )
                self.last_best_tag_id = None
                self.stable_count = 0
                return
            if self.max_correction_yaw_deg > 0.0 and correction_yaw > self.max_correction_yaw_deg:
                self.get_logger().warn(
                    f"Tag {observation.tag_id}: jump rejected - "
                    f"yaw {correction_yaw:.1f} deg > {self.max_correction_yaw_deg:.1f} deg; "
                    f"{self.localization_gate_summary()}"
                )
                self.last_best_tag_id = None
                self.stable_count = 0
                return

            blended = Pose2D(
                current.x + self.correction_alpha_pos * (candidate.x - current.x),
                current.y + self.correction_alpha_pos * (candidate.y - current.y),
                normalize_yaw(
                    current.yaw + self.correction_alpha_yaw * normalize_yaw(candidate.yaw - current.yaw)
                ),
            )
        else:
            blended = candidate

        self.get_logger().debug(
            f"Tag {observation.tag_id}: accepted "
            f"candidate=({candidate.x:.3f},{candidate.y:.3f},{math.degrees(candidate.yaw):.1f}deg) "
            f"current=({current.x:.3f},{current.y:.3f},{math.degrees(current.yaw):.1f}deg) "
            f"blended=({blended.x:.3f},{blended.y:.3f},{math.degrees(blended.yaw):.1f}deg) "
            f"correction=({correction_distance:.3f}m,{correction_yaw:.1f}deg) "
            f"yaw={observation.yaw_deg:.1f} pitch={observation.pitch_deg} roll={observation.roll_deg} "
            f"distance={observation.distance:.3f} reprojectionError={observation.reprojection_error} "
            f"cameraModel={observation.camera_model} poseImageSource={observation.pose_image_source}"
        )

        map_odom = map_odom_from_robot_pose(blended, odom_base_pose)
        self.set_map_odom(map_odom, f"tag:{observation.tag_id}", robot_map_yaw=blended.yaw)


def main() -> None:
    rclpy.init()
    node = TagLocalizationNode()
    try:
        rclpy.spin(node)
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == "__main__":
    main()
