import itertools
import random
import unittest

from dsu_lab import DisjointSet, road_repair


def connected(n, edges):
    adjacency = [set() for _ in range(n)]
    for a, b in edges:
        adjacency[a].add(b)
        adjacency[b].add(a)
    seen, stack = {0}, [0]
    while stack:
        for neighbor in adjacency[stack.pop()]:
            if neighbor not in seen:
                seen.add(neighbor)
                stack.append(neighbor)
    return len(seen) == n


def exhaustive_cost(n, existing, candidates):
    best = None
    for mask in range(1 << len(candidates)):
        chosen = [edge for i, edge in enumerate(candidates) if mask >> i & 1]
        if connected(n, existing + [(a, b) for a, b, _ in chosen]):
            cost = sum(price for _, _, price in chosen)
            best = cost if best is None else min(best, cost)
    return best


class RoadTests(unittest.TestCase):
    def test_worked_example(self):
        result = road_repair(5, [(0, 1), (1, 2), (0, 2)],
                             [(2, 3, 4), (1, 3, 8), (3, 4, 2), (0, 4, 10)])
        self.assertEqual(result["history"], [
            {"joined": True, "components": 4, "largest": 2},
            {"joined": True, "components": 3, "largest": 3},
            {"joined": False, "components": 3, "largest": 3}])
        self.assertEqual(result["repairs"], [(3, 4, 2), (2, 3, 4)])
        self.assertEqual((result["possible"], result["cost"], result["clusters"]),
                         (True, 6, [[0, 1, 2, 3, 4]]))

    def test_single_village(self):
        result = road_repair(1, [(0, 0)], [])
        self.assertTrue(result["possible"])
        self.assertEqual(result["cost"], 0)
        self.assertFalse(result["history"][0]["joined"])

    def test_already_connected_needs_no_purchase(self):
        result = road_repair(3, [(0, 1), (1, 2)], [(0, 2, 1)])
        self.assertEqual((result["cost"], result["repairs"]), (0, []))

    def test_disconnected_proposals(self):
        result = road_repair(4, [], [(0, 1, 2), (2, 3, 3)])
        self.assertFalse(result["possible"])
        self.assertIsNone(result["cost"])
        self.assertEqual(result["clusters"], [[0, 1], [2, 3]])

    def test_duplicate_and_zero_cost_roads(self):
        result = road_repair(3, [], [(1, 0, 0), (0, 1, 0), (1, 2, 0), (2, 2, 0)])
        self.assertEqual(result["repairs"], [(0, 1, 0), (1, 2, 0)])
        self.assertEqual(result["cost"], 0)

    def test_union_sizes_and_compression(self):
        groups = DisjointSet(8)
        for a, b in [(0, 1), (2, 3), (4, 5), (6, 7), (0, 2), (4, 6), (0, 4)]:
            self.assertTrue(groups.union(a, b))
        self.assertEqual(groups.components, 1)
        self.assertEqual(groups.size[groups.find(7)], 8)
        self.assertEqual(groups.find(7), groups.find(1))
        self.assertFalse(groups.union(7, 1))

    def test_invalid_input(self):
        for n, existing, candidates in [(0, [], []), (2, [(0, 2)], []),
                (2, [], [(0, 1, -1)]), (2, [], [(-1, 1, 2)])]:
            with self.subTest(n=n, candidates=candidates), self.assertRaises(ValueError):
                road_repair(n, existing, candidates)

    def test_random_cost_against_all_subsets(self):
        rng = random.Random(104)
        for _ in range(60):
            n = rng.randint(1, 5)
            pairs = list(itertools.combinations(range(n), 2))
            existing = rng.sample(pairs, rng.randrange(min(3, len(pairs)) + 1))
            candidates = [(a, b, rng.randrange(8)) for a, b in rng.sample(pairs, min(7, len(pairs)))]
            result = road_repair(n, existing, candidates)
            expected = exhaustive_cost(n, existing, candidates)
            self.assertEqual(result["cost"], expected)
            self.assertEqual(result["possible"], expected is not None)


if __name__ == "__main__":
    unittest.main()