"""Step 5 — 블록 분할 (MVP: partition만. balance_check/refine은 범위 밖)."""

from __future__ import annotations

import networkx as nx
import numpy as np
from sklearn.cluster import SpectralClustering

from .config import Config


def partition(Gw: nx.Graph, k: int, cfg: Config) -> dict[str, int]:
    if k < 1:
        raise ValueError(f"k는 1 이상이어야 합니다: {k}")
    if Gw.number_of_nodes() < k:
        raise ValueError(f"노드 수({Gw.number_of_nodes()})가 k({k})보다 적습니다")

    method = cfg.partition.method
    if method == "spectral":
        return _partition_spectral(Gw, k)
    if method == "louvain":
        return _partition_louvain(Gw, k)
    raise ValueError(f"알 수 없는 분할 방법: {method}")


def _partition_spectral(Gw: nx.Graph, k: int) -> dict[str, int]:
    nodes = list(Gw.nodes)
    A = nx.to_numpy_array(Gw, nodelist=nodes, weight="w")
    np.fill_diagonal(A, 0.0)

    model = SpectralClustering(
        n_clusters=k,
        affinity="precomputed",
        assign_labels="kmeans",
        random_state=0,
    )
    labels = model.fit_predict(A)
    return {n: int(label) for n, label in zip(nodes, labels)}


def _partition_louvain(Gw: nx.Graph, k: int, max_steps: int = 30) -> dict[str, int]:
    lo, hi = 0.01, 10.0
    best_communities: list[set] | None = None
    best_diff = None

    for _ in range(max_steps):
        r = (lo + hi) / 2
        communities = nx.community.louvain_communities(Gw, weight="w", resolution=r, seed=0)
        n_found = len(communities)
        diff = abs(n_found - k)
        if best_diff is None or diff < best_diff:
            best_diff, best_communities = diff, communities
        if n_found == k:
            break
        if n_found < k:
            lo = r
        else:
            hi = r

    labels: dict[str, int] = {}
    for i, community in enumerate(best_communities or []):
        for n in community:
            labels[n] = i
    return labels