"""map_exporter.py — Export RTAB-Map data to standard formats.

2D occupancy grid → PGM image + ROS map_server YAML
3D point cloud   → PLY or PCD file

Both exporters work by subscribing to RTAB-Map output topics and
capturing the latest message on demand.
"""

import struct
import time
import numpy as np
from datetime import datetime
from pathlib import Path


# ── 2D Occupancy Grid export ────────────────────────────────────────────────

def occupancy_grid_to_pgm(
    grid_data: list[int],
    width: int,
    height: int,
) -> bytes:
    """Convert ROS OccupancyGrid data to PGM (Portable GrayMap) bytes.

    OccupancyGrid encoding:
        -1  → unknown      → 205 (mid-grey in ROS convention)
         0  → free         → 254 (white)
        100 → occupied     → 0   (black)
        1-99 → partial     → interpolated

    Returns raw PGM P5 binary bytes.
    """
    pixels = np.zeros(len(grid_data), dtype=np.uint8)
    arr = np.array(grid_data, dtype=np.int8)

    unknown_mask  = arr == -1
    free_mask     = arr == 0
    occupied_mask = arr == 100
    partial_mask  = ~unknown_mask & ~free_mask & ~occupied_mask

    pixels[unknown_mask]  = 205
    pixels[free_mask]     = 254
    pixels[occupied_mask] = 0
    pixels[partial_mask]  = (254 - arr[partial_mask].astype(np.int32) * 254 // 100).astype(np.uint8)

    # PGM P5 (binary) format — rows stored bottom-up in ROS, flip vertically
    image = pixels.reshape(height, width)
    image = np.flipud(image)  # ROS maps have origin at bottom-left

    header = f'P5\n{width} {height}\n255\n'.encode()
    return header + image.tobytes()


def write_occupancy_export(
    grid_msg,            # nav_msgs/OccupancyGrid ROS message object
    output_dir: str,
    map_name: str = '',
) -> tuple[str, str]:
    """Write PGM + map_server YAML from an OccupancyGrid message.

    Returns:
        (pgm_path, yaml_path)
    """
    if not map_name:
        map_name = f'map_{datetime.now().strftime("%Y%m%d_%H%M%S")}'

    out = Path(output_dir)
    out.mkdir(parents=True, exist_ok=True)
    pgm_path  = out / f'{map_name}.pgm'
    yaml_path = out / f'{map_name}.yaml'

    info = grid_msg.info
    pgm_bytes = occupancy_grid_to_pgm(
        list(grid_msg.data), info.width, info.height)

    pgm_path.write_bytes(pgm_bytes)

    # map_server YAML format
    yaml_content = (
        f'image: {pgm_path.name}\n'
        f'resolution: {info.resolution:.4f}\n'
        f'origin: [{info.origin.position.x:.4f}, '
        f'{info.origin.position.y:.4f}, '
        f'{info.origin.position.z:.4f}]\n'
        f'negate: 0\n'
        f'occupied_thresh: 0.65\n'
        f'free_thresh: 0.196\n'
        f'mode: trinary\n'
    )
    yaml_path.write_text(yaml_content)

    return str(pgm_path), str(yaml_path)


# ── 3D Point Cloud export ────────────────────────────────────────────────────

def pointcloud2_to_xyz(cloud_msg) -> np.ndarray:
    """Extract XYZ points from a sensor_msgs/PointCloud2 message.

    Returns:
        np.ndarray of shape (N, 3) float32
    """
    # Parse field offsets
    fields = {f.name: f for f in cloud_msg.fields}
    x_off = fields['x'].offset
    y_off = fields['y'].offset
    z_off = fields['z'].offset
    point_step = cloud_msg.point_step
    data = bytes(cloud_msg.data)
    n_points = cloud_msg.width * cloud_msg.height

    xyz = np.zeros((n_points, 3), dtype=np.float32)
    for i in range(n_points):
        base = i * point_step
        xyz[i, 0] = struct.unpack_from('<f', data, base + x_off)[0]
        xyz[i, 1] = struct.unpack_from('<f', data, base + y_off)[0]
        xyz[i, 2] = struct.unpack_from('<f', data, base + z_off)[0]

    # Filter NaN/Inf
    valid = np.isfinite(xyz).all(axis=1)
    return xyz[valid]


def write_ply(xyz: np.ndarray, path: str) -> None:
    """Write XYZ point cloud as ASCII PLY file."""
    n = len(xyz)
    header = (
        'ply\n'
        'format ascii 1.0\n'
        f'element vertex {n}\n'
        'property float x\n'
        'property float y\n'
        'property float z\n'
        'end_header\n'
    )
    with open(path, 'w') as f:
        f.write(header)
        for p in xyz:
            f.write(f'{p[0]:.4f} {p[1]:.4f} {p[2]:.4f}\n')


def write_pcd(xyz: np.ndarray, path: str) -> None:
    """Write XYZ point cloud as ASCII PCD file (PCL-compatible)."""
    n = len(xyz)
    header = (
        '# .PCD v0.7 - Point Cloud Data file format\n'
        'VERSION 0.7\n'
        'FIELDS x y z\n'
        'SIZE 4 4 4\n'
        'TYPE F F F\n'
        'COUNT 1 1 1\n'
        f'WIDTH {n}\n'
        'HEIGHT 1\n'
        'VIEWPOINT 0 0 0 1 0 0 0\n'
        f'POINTS {n}\n'
        'DATA ascii\n'
    )
    with open(path, 'w') as f:
        f.write(header)
        for p in xyz:
            f.write(f'{p[0]:.4f} {p[1]:.4f} {p[2]:.4f}\n')


def write_pointcloud_export(
    cloud_msg,          # sensor_msgs/PointCloud2 ROS message object
    output_dir: str,
    map_name: str = '',
    fmt: str = 'ply',
) -> str:
    """Write point cloud to PLY or PCD.

    Returns:
        file_path
    """
    if fmt not in ('ply', 'pcd'):
        fmt = 'ply'
    if not map_name:
        map_name = f'cloud_{datetime.now().strftime("%Y%m%d_%H%M%S")}'

    out = Path(output_dir)
    out.mkdir(parents=True, exist_ok=True)
    file_path = out / f'{map_name}.{fmt}'

    xyz = pointcloud2_to_xyz(cloud_msg)
    if fmt == 'ply':
        write_ply(xyz, str(file_path))
    else:
        write_pcd(xyz, str(file_path))

    return str(file_path)
