"""Cluster a 2D LiDAR scan into world-frame sheep position estimates.

Pipeline:
    ranges (N,) ─► hit mask ─► world-frame points
                                     │
                                     ▼
                              adjacency clustering (gap > GAP_THRESHOLD
                              starts a new cluster, walking rays in
                              angular order)
                                     │
                                     ▼
                              centroid + span filter
                                     │
                                     ▼
                              field/pen-corridor filter
                                     │
                                     ▼
                              list of (x, y) detections

The clusterer is intentionally simple — for ≤10 sheep there is rarely
any real ambiguity, and proper DBSCAN would only matter if rays from
two adjacent sheep merged. The downstream tracker handles association
across frames.
"""

from __future__ import annotations

import math

import numpy as np

from herding.geometry import FIELD_X, FIELD_Y, GATE_Y, PEN_X, PEN_Y
from herding.lidar_sim import (
    LIDAR_FOV, LIDAR_MAX_RANGE, LIDAR_N_RAYS, SHEEP_RADIUS, ray_angles,
)


GAP_THRESHOLD = 0.6      # m — adjacent ray-points farther apart start new cluster
MAX_CLUSTER_SPAN = 1.5   # m — clusters wider than this are likely walls/structures
RANGE_HIT_EPS = 0.05     # m — hit if range < max_range - eps
WALL_REJECT = 0.5        # m — drop detections this close to a known wall line

# Known sheep-sized static features. Detections within STATIC_REJECT
# of any of these are discarded — these aren't sheep. Mid-pillars on
# the field walls are NOT in this list because they're embedded in the
# wall (the wall's span filter handles them); listing them here would
# only reject real sheep that happened to be near the wall.
_STATIC_FEATURES = (
    # Gate posts (sheep-sized boxes flanking the south-wall opening)
    ( 10.0, -15.0), ( 13.0, -15.0),
    # Field corner pillars
    ( 15.0,  15.0), ( 15.0, -15.0), (-15.0,  15.0), (-15.0, -15.0),
)
STATIC_REJECT = 0.8      # m — detection within this of a static feature → drop


def detections_from_scan(
    ranges: np.ndarray,
    dog_x: float, dog_y: float, dog_heading: float,
    max_range: float = LIDAR_MAX_RANGE,
) -> list[tuple[float, float]]:
    """Return list of (x, y) world-frame sheep position estimates."""
    ranges = np.asarray(ranges, dtype=np.float32)
    n_rays = ranges.shape[0]
    if n_rays == 0:
        return []
    angles = ray_angles(n_rays, LIDAR_FOV)
    hit = ranges < max_range - RANGE_HIT_EPS

    world_a = dog_heading + angles
    px = dog_x + ranges * np.cos(world_a)
    py = dog_y + ranges * np.sin(world_a)

    clusters: list[list[tuple[float, float]]] = []
    current: list[tuple[float, float]] = []
    prev: tuple[float, float] | None = None
    for i in range(n_rays):
        if not bool(hit[i]):
            if current:
                clusters.append(current)
                current = []
            prev = None
            continue
        pt = (float(px[i]), float(py[i]))
        if prev is not None and math.hypot(pt[0] - prev[0], pt[1] - prev[1]) > GAP_THRESHOLD:
            clusters.append(current)
            current = []
        current.append(pt)
        prev = pt
    if current:
        clusters.append(current)

    detections: list[tuple[float, float]] = []
    for cluster in clusters:
        xs = [p[0] for p in cluster]
        ys = [p[1] for p in cluster]
        cx, cy = sum(xs) / len(xs), sum(ys) / len(ys)
        span = math.hypot(max(xs) - min(xs), max(ys) - min(ys))
        if span > MAX_CLUSTER_SPAN:
            continue
        # Surface-to-centre correction: rays hit the front of the sheep,
        # so the cluster centroid is biased toward the dog by SHEEP_RADIUS.
        # Push it outward along the dog→cluster direction.
        dx, dy = cx - dog_x, cy - dog_y
        d = math.hypot(dx, dy)
        if d > 1e-3:
            cx += SHEEP_RADIUS * dx / d
            cy += SHEEP_RADIUS * dy / d
        # Keep detections inside the field OR in the gate corridor /
        # external pen — penned sheep are still worth tracking so the
        # tracker can latch them as "penned" rather than spawn fresh
        # tracks each scan.
        # Accept detections inside the field, plus a narrow strip
        # immediately south of the gate to catch sheep mid-crossing
        # (so they get marked penned via is_penned_position before the
        # track goes stale). Detections deeper into the pen are
        # dropped entirely — Webots's pen posts and rails would
        # otherwise produce a torrent of phantom penned tracks that
        # the tracker can't keep up with.
        in_main = (FIELD_X[0] - 0.2 < cx < FIELD_X[1] + 0.2 and
                   FIELD_Y[0] - 0.2 < cy < FIELD_Y[1] + 0.2)
        in_gate_strip = (PEN_X[0] - 0.2 < cx < PEN_X[1] + 0.2 and
                         GATE_Y - 1.0 < cy < GATE_Y + 0.2)
        if not (in_main or in_gate_strip):
            continue
        # Known-static-feature filter: gate posts and corner pillars
        # show up as sheep-sized clusters but are never sheep.
        if any(math.hypot(cx - fx, cy - fy) < STATIC_REJECT
               for fx, fy in _STATIC_FEATURES):
            continue
        # Wall-proximity filter: at oblique scan angles, walls produce
        # multiple short clusters because adjacent ray returns are
        # spaced just above GAP_THRESHOLD. Sheep can't get within ~0.3 m
        # of a wall (the env clips them to FIELD_INSIDE), so anything
        # right at the wall line is structure noise.
        near_field_wall = (
            cx > FIELD_X[1] - WALL_REJECT or cx < FIELD_X[0] + WALL_REJECT or
            cy > FIELD_Y[1] - WALL_REJECT or
            (cy < FIELD_Y[0] + WALL_REJECT and not (PEN_X[0] <= cx <= PEN_X[1]))
        )
        if near_field_wall:
            continue
        detections.append((cx, cy))
    return detections