"""
map_compressor.py — Serialise / compress / decompress OccupancyGrid for fleet transport.

OccupancyGrid serialisation format (little-endian binary):
  4B  resolution (float32)
  4B  origin_x   (float32)
  4B  origin_y   (float32)
  4B  origin_yaw (float32)   — extracted from quaternion
  4B  width      (uint32)
  4B  height     (uint32)
  N   data       (int8 × width*height)

All compressed with zlib level-6.  Decompression restores a nav_msgs/OccupancyGrid
with the same metadata so the receiver can directly publish / merge it.
"""

import struct
import zlib
import math
from typing import Tuple

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


_HEADER_FMT = '<ffffII'   # resolution, ox, oy, oyaw, width, height
_HEADER_SIZE = struct.calcsize(_HEADER_FMT)


def _quat_to_yaw(q: Quaternion) -> float:
    siny = 2.0 * (q.w * q.z + q.x * q.y)
    cosy = 1.0 - 2.0 * (q.y * q.y + q.z * q.z)
    return math.atan2(siny, cosy)


def _yaw_to_quat(yaw: float) -> Quaternion:
    q = Quaternion()
    q.w = math.cos(yaw * 0.5)
    q.z = math.sin(yaw * 0.5)
    return q


def compress_grid(grid: OccupancyGrid) -> Tuple[bytes, int]:
    """
    Serialise and zlib-compress an OccupancyGrid.

    Returns:
        (compressed_bytes, uncompressed_size)
    """
    info = grid.info
    yaw = _quat_to_yaw(info.origin.orientation)
    header = struct.pack(
        _HEADER_FMT,
        info.resolution,
        info.origin.position.x,
        info.origin.position.y,
        yaw,
        info.width,
        info.height,
    )
    data = bytes([v & 0xFF for v in grid.data])   # int8 → uint8 for bytes()
    raw = header + data
    return zlib.compress(raw, level=6), len(raw)


def decompress_grid(compressed: bytes, frame_id: str = 'map') -> OccupancyGrid:
    """
    Decompress and deserialise bytes produced by compress_grid() into an OccupancyGrid.
    """
    raw = zlib.decompress(compressed)
    resolution, ox, oy, oyaw, width, height = struct.unpack_from(_HEADER_FMT, raw, 0)
    data_bytes = raw[_HEADER_SIZE:]

    grid = OccupancyGrid()
    grid.header.frame_id = frame_id
    grid.info.resolution = resolution
    grid.info.width = width
    grid.info.height = height
    grid.info.origin = Pose(
        position=Point(x=float(ox), y=float(oy), z=0.0),
        orientation=_yaw_to_quat(oyaw),
    )
    # Re-sign: stored as uint8, OccupancyGrid.data is int8 list
    grid.data = [int(b) if b < 128 else int(b) - 256 for b in data_bytes]
    return grid


def compression_ratio(compressed: bytes, uncompressed_size: int) -> float:
    return len(compressed) / max(uncompressed_size, 1)
