"""
landmark_aligner.py — Estimate SE(2) transforms between robot map frames
using common landmark observations.

Each robot's RTAB-Map publishes loop-closure keyframes with GPS-less
"global" constraints derived from visual features.  When two robots
independently observe the same physical landmark (ArUco marker, stable
visual feature cluster), their descriptor vectors will match and we can
compute the rigid transform that maps robot-A's map frame to robot-B's
(or the common fleet_map frame).

Approach used here:
  - Robots publish detected ArUco marker poses on /<ns>/landmarks
    (geometry_msgs/PoseArray, frame = <ns>/map).
  - The LandmarkAligner collects observations per marker_id from all robots.
  - When the same marker_id has been seen by ≥ 2 robots, solve for the
    SE(2) transform between each robot's map frame and the canonical
    fleet_map frame (defined as robot_1's map frame by convention).
  - Transform is published as a geometry_msgs/TransformStamped and also
    returned as a geometry_msgs/Transform for embedding in MapChunk.

Landmark topic convention:
  /<robot_ns>/landmarks      PoseArray
    - poses[i].position.x/y   = landmark centroid in robot's map frame
    - poses[i].position.z     = float-encoded marker_id (cast to int)

SE(2) estimation:
  Uses the Umeyama / Horn algorithm on 2-D point correspondences.
  Rotation + translation only (no scale).

Module is standalone — no ROS node, just math helpers called from
fleet_manager_node.py.
"""

from __future__ import annotations
import math
from typing import Dict, List, Optional, Tuple

import numpy as np


# robot_id → {marker_id → (x, y) in that robot's map frame}
_LandmarkObs = Dict[str, Dict[int, Tuple[float, float]]]


def _se2_from_correspondences(
    src_pts: np.ndarray,   # (N, 2) points in source frame
    dst_pts: np.ndarray,   # (N, 2) points in destination frame
) -> Tuple[float, float, float]:
    """
    Compute SE(2) transform (tx, ty, yaw) that maps src_pts → dst_pts.

    Uses the closed-form Horn least-squares solution for 2-D rigid body.
    Returns (tx, ty, yaw).
    """
    assert src_pts.shape == dst_pts.shape and src_pts.shape[1] == 2
    n = len(src_pts)
    if n < 2:
        raise ValueError('Need ≥ 2 correspondences for SE(2) estimation')

    # Centre the point sets
    c_src = src_pts.mean(axis=0)
    c_dst = dst_pts.mean(axis=0)
    s     = src_pts - c_src
    d     = dst_pts - c_dst

    # Cross-covariance  H = S^T D
    H = s.T @ d  # (2, 2)
    # Rotation via atan2 of the off-diagonal elements of H
    yaw = math.atan2(H[0, 1] - H[1, 0], H[0, 0] + H[1, 1])
    c, ss = math.cos(yaw), math.sin(yaw)
    R = np.array([[c, -ss], [ss, c]])
    t = c_dst - R @ c_src
    return float(t[0]), float(t[1]), float(yaw)


def _build_transform(tx: float, ty: float, yaw: float):
    """Return geometry_msgs/Transform-like dict."""
    from geometry_msgs.msg import Transform, Vector3, Quaternion
    tf = Transform()
    tf.translation.x = tx
    tf.translation.y = ty
    tf.translation.z = 0.0
    tf.rotation.w = math.cos(yaw * 0.5)
    tf.rotation.z = math.sin(yaw * 0.5)
    return tf


class LandmarkAligner:
    """
    Collects landmark observations from multiple robots and estimates
    the per-robot transform into the fleet_map frame.

    fleet_map is defined as robot_1's map frame (robot_id with the
    lowest registration sequence number in the coordinator).
    """

    def __init__(self, reference_robot: str):
        """
        Args:
            reference_robot: robot_id whose map frame IS fleet_map.
                             Other robots are aligned to this reference.
        """
        self._reference   = reference_robot
        self._obs: _LandmarkObs = {}   # robot_id → {marker_id → (x, y)}
        # robot_id → (tx, ty, yaw, valid)
        self._transforms: Dict[str, Tuple[float, float, float, bool]] = {}

    def update_observations(
        self,
        robot_id: str,
        landmarks: List[Tuple[int, float, float]],  # [(marker_id, x, y), ...]
    ) -> None:
        """Record landmark observations from a robot."""
        if robot_id not in self._obs:
            self._obs[robot_id] = {}
        for mid, x, y in landmarks:
            self._obs[robot_id][mid] = (x, y)
        self._recompute(robot_id)

    def _recompute(self, updated_robot: str) -> None:
        """Recompute transform for updated_robot against the reference."""
        if updated_robot == self._reference:
            self._transforms[updated_robot] = (0.0, 0.0, 0.0, True)
            return
        ref_obs = self._obs.get(self._reference)
        rob_obs = self._obs.get(updated_robot)
        if not ref_obs or not rob_obs:
            return
        # Find shared marker IDs
        common = set(ref_obs.keys()) & set(rob_obs.keys())
        if len(common) < 2:
            return  # insufficient correspondences
        src = np.array([rob_obs[mid] for mid in common])
        dst = np.array([ref_obs[mid] for mid in common])
        try:
            tx, ty, yaw = _se2_from_correspondences(src, dst)
        except ValueError:
            return
        self._transforms[updated_robot] = (tx, ty, yaw, True)

    def get_transform(self, robot_id: str):
        """
        Return geometry_msgs/Transform for robot_id's map → fleet_map.
        Returns (transform, valid:bool).
        """
        entry = self._transforms.get(robot_id)
        if entry is None or not entry[3]:
            from geometry_msgs.msg import Transform
            return Transform(), False
        tx, ty, yaw, _ = entry
        return _build_transform(tx, ty, yaw), True

    def aligned_robots(self) -> List[str]:
        return [rid for rid, t in self._transforms.items() if t[3]]
