"""Cluster detection for AFS colony results."""

import numpy as np
from dataclasses import dataclass


@dataclass
class Cluster:
    centroid: np.ndarray
    cell_count: int
    member_indices: list[int]


def detect_clusters(cell_positions: np.ndarray, max_clusters: int = 8) -> list[Cluster]:
    """
    Density-based cluster detection.

    Same algorithm as JS AfsCompiler._detectClusters, but vectorized with NumPy.
    Finds the densest unassigned cell, collects neighbors within radius, repeats.

    Args:
        cell_positions: (N, D) array of cell positions
        max_clusters: maximum clusters to detect

    Returns:
        List of Cluster objects sorted by cell count (descending)
    """
    n_cells, n_dims = cell_positions.shape
    if n_cells == 0:
        return []

    # Scale radius sub-linearly with dimensions to avoid swallowing
    # the entire space into one cluster in high-D.
    # sqrt grows too fast: 10D → 0.47 (half the space).
    # log grows right:     10D → 0.19, 4D → 0.14, 15D → 0.22.
    cluster_radius = 0.10 + 0.04 * np.log(n_dims)
    min_cluster_size = max(3, int(n_cells * 0.05))
    assigned = np.zeros(n_cells, dtype=bool)
    clusters = []

    # Precompute pairwise distances
    dists = np.linalg.norm(
        cell_positions[:, None, :] - cell_positions[None, :, :], axis=2
    )  # (N, N)

    while np.sum(~assigned) > min_cluster_size and len(clusters) < max_clusters:
        # Count neighbors for each unassigned cell
        unassigned_mask = ~assigned
        neighbor_counts = np.sum(
            (dists < cluster_radius) & unassigned_mask[None, :], axis=1
        )
        # Zero out assigned cells so they can't be seeds
        neighbor_counts[assigned] = 0

        best_idx = int(np.argmax(neighbor_counts))
        best_count = int(neighbor_counts[best_idx])

        if best_count < min_cluster_size:
            break

        # Collect members: unassigned cells within radius of best
        member_mask = (dists[best_idx] < cluster_radius) & unassigned_mask
        member_indices = np.where(member_mask)[0].tolist()
        assigned[member_mask] = True

        centroid = np.mean(cell_positions[member_mask], axis=0)
        clusters.append(Cluster(
            centroid=centroid,
            cell_count=len(member_indices),
            member_indices=member_indices,
        ))

    clusters.sort(key=lambda c: c.cell_count, reverse=True)
    return clusters


def compute_concept_weights(
    cluster: Cluster,
    concept_positions: np.ndarray,
    concept_names: list[str],
    n_dims: int,
) -> dict[str, float]:
    """
    Compute which concepts dominate a cluster.

    Uses inverse distance weighting from cluster centroid to each concept position.
    """
    weights = {}
    for i, name in enumerate(concept_names):
        dist = float(np.linalg.norm(cluster.centroid[:n_dims] - concept_positions[i, :n_dims]))
        weight = 1.0 / (1.0 + dist * 5.0)
        weights[name] = round(weight, 3)
    return weights
