from __future__ import annotations

import math

from pygui.backend.wafer.zwafer_models import Sensor


def ring_boundaries(sensors: list[Sensor]) -> list[float]:
    """Distinct radial distances of sensor groups, sorted ascending, plus the outer boundary."""
    if not sensors:
        r = 150.0
        return [r * f for f in (0.25, 0.50, 0.75, 1.0)]
    # Cluster radii that are within 2 mm of each other into one ring; skip center point.
    radii = sorted(r for r in {math.hypot(s.x, s.y) for s in sensors} if r > 1.0)
    groups: list[float] = []
    for r in radii:
        if not groups or r - groups[-1] > 2.0:
            groups.append(r)
        else:
            groups[-1] = (groups[-1] + r) / 2  # merge close values
    # Always include the outer boundary ring so the wafer circle is drawn.
    outer = max(math.hypot(s.x, s.y) for s in sensors) * 1.05
    if not groups or outer - groups[-1] > 2.0:
        groups.append(outer)
    return groups


def bucket_by_ring(sensors: list[Sensor], values: list[float]) -> dict[int, list[float]]:
    """Group each sensor's value under the index of its nearest ring boundary."""
    boundaries = ring_boundaries(sensors)
    result: dict[int, list[float]] = {i: [] for i in range(len(boundaries))}
    for i, s in enumerate(sensors):
        if i >= len(values):
            continue
        r = math.hypot(s.x, s.y)
        idx = min(range(len(boundaries)), key=lambda j: abs(boundaries[j] - r))
        result[idx].append(values[i])
    return result
