import itertools
import random
import unittest
from example import DisjointSet, minimum_spanning_forest


class MSTTests(unittest.TestCase):
    def test_forest_negative_self_loop(self):
        cost, chosen, components = minimum_spanning_forest(4, [(-9, 0, 0), (-2, 0, 1), (3, 0, 1)])
        self.assertEqual((cost, len(chosen), components), (-2, 1, 3))
        self.assertEqual(minimum_spanning_forest(0, []), (0, [], 0))

    def test_against_exhaustive_connected_graphs(self):
        rng = random.Random(15)
        for _ in range(25):
            n = 4
            edges = [(rng.randint(-3, 8), a, b) for a in range(n) for b in range(a + 1, n)]
            best = float('inf')
            for candidate in itertools.combinations(edges, n - 1):
                ds = DisjointSet(n)
                if all(ds.union(a, b) for _, a, b in candidate):
                    best = min(best, sum(w for w, _, _ in candidate))
            cost, chosen, components = minimum_spanning_forest(n, edges)
            self.assertEqual(cost, best)
            self.assertEqual(components, 1)
            self.assertEqual(len(chosen), n - 1)

    def test_invalid_vertex(self):
        with self.assertRaises(ValueError):
            minimum_spanning_forest(2, [(1, 0, 2)])


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