"""
map_merger.py — Merge multiple OccupancyGrids into a single fleet-wide map.

Each robot publishes its local occupancy grid compressed on /fleet/maps.
The coordinator decompresses them and calls merge_grids() to produce a unified
merged map suitable for frontier detection and visualisation on /fleet/merged_map.

Merge strategy:
  - Compute the bounding box that covers all incoming grids (same resolution required).
  - Accumulate probability estimates per cell: average known values; unknown (-1) only
    fills cells not covered by any robot.
  - Occupied (> 65) beats free (0–65) when two robots disagree on a cell — conservative
    obstacle inflation to prevent Nav2 from driving into conflicting-report zones.

Transform alignment:
  - Each MapChunk carries a map_to_fleet_tf transform (set by the coordinator after
    landmark matching).  If tf_valid=true the grid is rotated/translated into the
    common fleet frame before merging.
  - Until tf_valid is set, the grid is assumed already in the common frame (works
    when all robots start from the same origin or when only one robot is active).

Resolution constraint:
  - All grids MUST share the same resolution.  If they differ, the chunk is discarded
    and a warning is logged.  Re-parameterise RTAB-Map Grid/CellSize to enforce 0.05m.
"""

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

import numpy as np

from nav_msgs.msg import OccupancyGrid
from geometry_msgs.msg import Pose, Point, Quaternion, Transform


_FREE     =  0
_OCCUPIED = 100
_UNKNOWN  = -1


def _transform_to_matrix(tf: Transform) -> np.ndarray:
    """Return 3×3 SE(2) homogeneous matrix from a geometry_msgs/Transform."""
    tx = tf.translation.x
    ty = tf.translation.y
    q  = tf.rotation
    yaw = math.atan2(
        2.0 * (q.w * q.z + q.x * q.y),
        1.0 - 2.0 * (q.y * q.y + q.z * q.z),
    )
    c, s = math.cos(yaw), math.sin(yaw)
    return np.array([
        [c, -s, tx],
        [s,  c, ty],
        [0,  0,  1],
    ])


def _grid_bounds(grid: OccupancyGrid, tf_mat: np.ndarray) -> Tuple[float, float, float, float]:
    """Return (min_x, min_y, max_x, max_y) of grid corners after transform."""
    info = grid.info
    ox, oy = info.origin.position.x, info.origin.position.y
    w, h   = info.width * info.resolution, info.height * info.resolution
    corners = np.array([
        [ox,     oy,     1],
        [ox + w, oy,     1],
        [ox + w, oy + h, 1],
        [ox,     oy + h, 1],
    ]).T                          # shape (3, 4)
    transformed = tf_mat @ corners    # shape (3, 4)
    xs, ys = transformed[0], transformed[1]
    return float(xs.min()), float(ys.min()), float(xs.max()), float(ys.max())


class MapMerger:
    """
    Accumulates OccupancyGrids from multiple robots and produces a merged map.

    Usage:
        merger = MapMerger(resolution=0.05)
        merger.update('robot_1', grid1, tf_valid=False)
        merger.update('robot_2', grid2, tf_valid=True, tf=some_transform)
        merged = merger.build_merged_grid()
    """

    def __init__(self, resolution: float = 0.05):
        self._resolution   = resolution
        # robot_id → (grid: OccupancyGrid, tf_matrix: np.ndarray)
        self._grids: Dict[str, Tuple[OccupancyGrid, np.ndarray]] = {}

    def update(
        self,
        robot_id: str,
        grid: OccupancyGrid,
        tf_valid: bool = False,
        tf: Transform | None = None,
    ) -> bool:
        """
        Register or update a robot's latest grid.  Returns False if resolution mismatch.
        """
        if abs(grid.info.resolution - self._resolution) > 1e-5:
            return False
        tf_mat = _transform_to_matrix(tf) if (tf_valid and tf is not None) else np.eye(3)
        self._grids[robot_id] = (grid, tf_mat)
        return True

    def remove(self, robot_id: str) -> None:
        self._grids.pop(robot_id, None)

    def build_merged_grid(self) -> OccupancyGrid | None:
        """
        Merge all registered grids into a single OccupancyGrid.
        Returns None if no grids are registered.
        """
        if not self._grids:
            return None

        res = self._resolution

        # ── Compute merged bounding box ────────────────────────────────────
        min_x = min_y =  math.inf
        max_x = max_y = -math.inf
        for grid, tf_mat in self._grids.values():
            bx0, by0, bx1, by1 = _grid_bounds(grid, tf_mat)
            min_x = min(min_x, bx0)
            min_y = min(min_y, by0)
            max_x = max(max_x, bx1)
            max_y = max(max_y, by1)

        W = int(math.ceil((max_x - min_x) / res)) + 2
        H = int(math.ceil((max_y - min_y) / res)) + 2
        # Accumulators: sum of known values, count of known observations
        acc   = np.full((H, W), 0.0, dtype=np.float32)
        count = np.zeros((H, W), dtype=np.int32)

        # ── Paint each robot's grid into the merged canvas ─────────────────
        for grid, tf_mat in self._grids.values():
            src = np.array(grid.data, dtype=np.int8).reshape(
                grid.info.height, grid.info.width
            )
            ox = grid.info.origin.position.x
            oy = grid.info.origin.position.y

            rows, cols = np.where(src != _UNKNOWN)
            values = src[rows, cols].astype(np.float32)

            # Cell centre in source grid frame
            px = ox + (cols + 0.5) * res
            py = oy + (rows + 0.5) * res
            ones = np.ones_like(px)
            pts = np.stack([px, py, ones], axis=0)          # (3, N)
            transformed = tf_mat @ pts                       # (3, N)

            # Target cell indices in merged canvas
            tcol = ((transformed[0] - min_x) / res).astype(int)
            trow = ((transformed[1] - min_y) / res).astype(int)

            valid = (tcol >= 0) & (tcol < W) & (trow >= 0) & (trow < H)
            tcol, trow, values = tcol[valid], trow[valid], values[valid]

            # Conservative merge: occupied wins over free in conflicts
            existing_occ = (acc[trow, tcol] > 65) & (count[trow, tcol] > 0)
            incoming_occ = values > 65
            # For already-occupied cells don't average down with free
            update_mask = ~existing_occ | incoming_occ
            acc[trow[update_mask], tcol[update_mask]]   += values[update_mask]
            count[trow[update_mask], tcol[update_mask]] += 1

        # ── Build output OccupancyGrid ─────────────────────────────────────
        merged = OccupancyGrid()
        merged.header.frame_id = 'fleet_map'
        merged.info.resolution = res
        merged.info.width  = W
        merged.info.height = H
        merged.info.origin.position.x = min_x
        merged.info.origin.position.y = min_y
        merged.info.origin.orientation.w = 1.0

        merged_arr = np.full((H, W), _UNKNOWN, dtype=np.int8)
        known = count > 0
        merged_arr[known] = np.clip(
            (acc[known] / count[known]).astype(np.int8), 0, 100
        )
        merged.data = merged_arr.flatten().tolist()
        return merged
