import random
from dsu import DSU


def test_basic():
    d = DSU(5)
    assert d.components == 5
    assert not d.connected(0, 1)
    assert d.union(0, 1) is True
    assert d.connected(0, 1)
    assert d.components == 4
    assert d.union(0, 1) is False
    assert d.components == 4


def test_size():
    d = DSU(6)
    d.union(0, 1); d.union(1, 2)
    assert d.size(0) == 3 and d.size(2) == 3
    assert d.size(3) == 1
    d.union(3, 4)
    assert d.size(3) == 2


def test_transitive():
    d = DSU(4)
    d.union(0, 1); d.union(2, 3); d.union(1, 2)
    assert d.components == 1
    for i in range(4):
        for j in range(4):
            assert d.connected(i, j)


def test_empty():
    d = DSU(0)
    assert d.components == 0
    assert len(d) == 0
    assert d.sets() == {}


def test_reset():
    d = DSU(4)
    d.union(0, 1); d.union(2, 3)
    d.reset()
    assert d.components == 4
    assert not d.connected(0, 1)


def test_sets():
    d = DSU(5)
    d.union(0, 2); d.union(1, 3)
    s = d.sets()
    assert sorted(s.values(), key=lambda v: v[0]) == [[0, 2], [1, 3], [4]]


def test_errors():
    d = DSU(3)
    try:
        d.find(5)
    except IndexError:
        pass
    else:
        raise AssertionError("expected IndexError")
    try:
        DSU(-1)
    except ValueError:
        pass
    else:
        raise AssertionError("expected ValueError")
    try:
        DSU(True)
    except ValueError:
        pass
    else:
        raise AssertionError("expected ValueError for bool")


def test_random_vs_reference():
    rng = random.Random(42)
    n = 30
    d = DSU(n)
    parent = list(range(n))
    def ref_find(x):
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x
    for _ in range(2000):
        a = rng.randrange(n); b = rng.randrange(n)
        # union in both
        ra, rb = ref_find(a), ref_find(b)
        expected_merge = ra != rb
        if expected_merge:
            parent[ra] = rb
        got = d.union(a, b)
        assert got == expected_merge, (a, b, got, expected_merge)
        # connectivity check sample
        for _ in range(5):
            x = rng.randrange(n); y = rng.randrange(n)
            assert d.connected(x, y) == (ref_find(x) == ref_find(y))
    assert d.components == len({ref_find(i) for i in range(n)})


def test_deep_chain_no_recursion():
    # union in a line to force deep chains, then find must not recurse
    n = 100000
    d = DSU(n)
    for i in range(n - 1):
        d.union(i, i + 1)
    assert d.components == 1
    assert d.find(0) == d.find(n - 1)
    assert d.size(0) == n