from __future__ import annotations

import networkx as nx
import pytest

from blocksys.partition import partition

K_VALUES = [4, 6, 8]


@pytest.mark.parametrize("k", K_VALUES)
def test_node_preservation(weighted_graph, cfg, k):
    labels = partition(weighted_graph, k, cfg)
    assert set(labels.keys()) == set(weighted_graph.nodes)


@pytest.mark.parametrize("k", K_VALUES)
def test_block_connectivity(weighted_graph, cfg, k):
    labels = partition(weighted_graph, k, cfg)
    blocks: dict[int, list] = {}
    for n, label in labels.items():
        blocks.setdefault(label, []).append(n)

    for label, nodes in blocks.items():
        sub = weighted_graph.subgraph(nodes)
        assert nx.is_connected(sub), f"block {label} is not a single connected component"


def test_reproducibility(weighted_graph, cfg):
    labels1 = partition(weighted_graph, 6, cfg)
    labels2 = partition(weighted_graph, 6, cfg)
    assert labels1 == labels2


def test_k_too_large_raises(weighted_graph, cfg):
    with pytest.raises(ValueError):
        partition(weighted_graph, weighted_graph.number_of_nodes() + 1, cfg)


def test_louvain_method(weighted_graph, cfg):
    cfg.partition.method = "louvain"
    try:
        labels = partition(weighted_graph, 6, cfg)
        assert set(labels.keys()) == set(weighted_graph.nodes)
    finally:
        cfg.partition.method = "spectral"