"""Cluster Average utility for wafer sensor values."""


from __future__ import annotations

import math

from pygui.backend.wafer.zwafer_models import Sensor


def average_clusters(values: list[float], clusters: list[list[int]]) -> list[float]:
    """Return a new list of values where members of each cluster are replaced by their mean.

    NaN values are filtered out before computing the mean.
    If all values in a cluster are NaN, they remain NaN.
    """
    result = list(values)
    for cluster in clusters:
        valid_vals = []
        for idx in cluster:
            if idx < len(values):
                val = values[idx]
                if not math.isnan(val):
                    valid_vals.append(val)

        if valid_vals:
            mean_val = sum(valid_vals) / len(valid_vals)
            for idx in cluster:
                if idx < len(result):
                    result[idx] = mean_val
    return result


def group_sensors_by_radius(sensors: list[Sensor], tolerance: float = 2.0) -> list[list[int]]:
    """Group sensor indices into clusters if their radial distances are within tolerance."""
    # Compute radii for all sensors
    radii = [(i, math.hypot(s.x, s.y)) for i, s in enumerate(sensors)]
    # Filter out center sensors (r < 1.0)
    non_center = [item for item in radii if item[1] >= 1.0]
    if not non_center:
        return []

    # Sort by radius
    non_center.sort(key=lambda x: x[1])

    clusters: list[list[int]] = []
    current_cluster: list[int] = [non_center[0][0]]
    last_r = non_center[0][1]

    for idx, r in non_center[1:]:
        if r - last_r <= tolerance:
            current_cluster.append(idx)
        else:
            if len(current_cluster) > 1:
                clusters.append(current_cluster)
            current_cluster = [idx]
            last_r = r

    if len(current_cluster) > 1:
        clusters.append(current_cluster)

    return clusters