"""
obstacle_clusterer.py — 2D grid-based obstacle clustering (Issue #611).

No ROS2 dependencies — pure Python + numpy, fully unit-testable.

Algorithm
─────────
After RANSAC ground removal, obstacle points (non-ground) are clustered
using a 2D occupancy grid on the horizontal (X–Z) plane in camera frame.

Steps:
  1. Project obstacle 3D points onto the (X, Z) bird's-eye plane.
  2. Discretise into grid cells of side `cell_m`.
  3. 4-connected BFS flood-fill to find connected components.
  4. For each component with ≥ min_pts 3D points:
       centroid = mean of the component's 3D points
       radius   = max distance from centroid (bounding-sphere radius)

Camera frame convention:
  +X = right, +Y = down (not used for 2D projection), +Z = forward (depth)

Output per cluster
──────────────────
  ObstacleCluster (dataclass):
    centroid  : (3,) float array  — (cx, cy, cz) in camera frame
    radius_m  : float             — bounding-sphere radius (m)
    height_m  : float             — mean height above ground plane (m)
    n_pts     : int               — number of 3D points in cluster
"""

from __future__ import annotations

from collections import deque
from dataclasses import dataclass, field
from typing import List, Tuple

import numpy as np


@dataclass
class ObstacleCluster:
    centroid: np.ndarray          # (3,) float64 in camera optical frame
    radius_m: float
    height_m: float
    n_pts:    int

    @property
    def distance_m(self) -> float:
        """Forward distance (Z in camera frame)."""
        return float(self.centroid[2])

    @property
    def lateral_m(self) -> float:
        """Lateral offset (X in camera frame; positive = right)."""
        return float(self.centroid[0])


def cluster_obstacles(
    points:    np.ndarray,
    heights:   np.ndarray,
    cell_m:    float = 0.30,
    min_pts:   int   = 5,
    x_range:   Tuple[float, float] = (-4.0, 4.0),
    z_range:   Tuple[float, float] = (0.15, 6.0),
) -> List[ObstacleCluster]:
    """
    Cluster obstacle points into connected groups.

    Parameters
    ----------
    points   : (N, 3) float array — obstacle points in camera frame
    heights  : (N,)  float array — height above ground for each point
    cell_m   : grid cell side length (m)
    min_pts  : minimum points for a cluster to be reported
    x_range  : (min_x, max_x) bounds for the bird's-eye grid (m)
    z_range  : (min_z, max_z) forward-distance bounds (m)

    Returns
    -------
    clusters : list of ObstacleCluster, sorted by ascending distance (z)
    """
    if len(points) == 0:
        return []

    pts = points.astype(np.float64)

    # ── Discard out-of-range points ─────────────────────────────────────────
    x = pts[:, 0]
    z = pts[:, 2]
    valid = (
        (x >= x_range[0]) & (x <= x_range[1]) &
        (z >= z_range[0]) & (z <= z_range[1])
    )
    pts     = pts[valid]
    heights = heights.astype(np.float64)[valid]

    if len(pts) == 0:
        return []

    # ── Build 2D occupancy grid (X, Z) ──────────────────────────────────────
    x_min, x_max = x_range
    z_min, z_max = z_range
    n_x = max(1, int(math.ceil((x_max - x_min) / cell_m)))
    n_z = max(1, int(math.ceil((z_max - z_min) / cell_m)))

    xi = np.clip(((pts[:, 0] - x_min) / cell_m).astype(np.int32), 0, n_x - 1)
    zi = np.clip(((pts[:, 2] - z_min) / cell_m).astype(np.int32), 0, n_z - 1)

    # Map each grid cell to list of 3D-point indices
    cell_pts: dict[Tuple[int, int], list] = {}
    for idx, (gx, gz) in enumerate(zip(xi.tolist(), zi.tolist())):
        key = (gx, gz)
        if key not in cell_pts:
            cell_pts[key] = []
        cell_pts[key].append(idx)

    occupied = set(cell_pts.keys())

    # ── BFS connected-component labelling ────────────────────────────────────
    visited  = set()
    clusters = []

    for start in occupied:
        if start in visited:
            continue

        # BFS flood fill
        component_cells = []
        queue = deque([start])
        visited.add(start)

        while queue:
            cx, cz = queue.popleft()
            component_cells.append((cx, cz))

            for dx, dz in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
                nb = (cx + dx, cz + dz)
                if nb in occupied and nb not in visited:
                    visited.add(nb)
                    queue.append(nb)

        # Collect 3D points in this component
        pt_indices = []
        for cell in component_cells:
            pt_indices.extend(cell_pts[cell])

        if len(pt_indices) < min_pts:
            continue

        cluster_pts    = pts[pt_indices]
        cluster_h      = heights[pt_indices]
        centroid       = cluster_pts.mean(axis=0)
        dists          = np.linalg.norm(cluster_pts - centroid, axis=1)
        radius_m       = float(dists.max())
        mean_height    = float(cluster_h.mean())

        clusters.append(ObstacleCluster(
            centroid = centroid,
            radius_m = radius_m,
            height_m = mean_height,
            n_pts    = len(pt_indices),
        ))

    # Sort by forward distance
    clusters.sort(key=lambda c: c.distance_m)
    return clusters


# ── Local import for math ─────────────────────────────────────────────────────
import math
