import sys, os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import unittest
from dsu_rollback import DSU


class TestDSURollback(unittest.TestCase):
    def test_basic_union_connected(self):
        d = DSU(4)
        self.assertTrue(d.union(0, 1))
        self.assertTrue(d.connected(0, 1))
        self.assertFalse(d.connected(0, 2))
        self.assertEqual(d.size(0), 2)

    def test_noop_union_returns_false_no_history(self):
        d = DSU(3)
        d.union(0, 1)
        m = d.snapshot()
        self.assertFalse(d.union(0, 1))
        self.assertEqual(d.snapshot(), m)

    def test_rollback_restores(self):
        d = DSU(5)
        d.union(0, 1); d.union(2, 3)
        m = d.snapshot()
        d.union(1, 2); d.union(3, 4)
        self.assertTrue(d.connected(0, 4))
        d.rollback(m)
        self.assertFalse(d.connected(0, 4))
        self.assertTrue(d.connected(0, 1))
        self.assertTrue(d.connected(2, 3))
        self.assertFalse(d.connected(1, 2))

    def test_nested_snapshots(self):
        d = DSU(6)
        d.union(0, 1)
        m1 = d.snapshot()
        d.union(2, 3)
        m2 = d.snapshot()
        d.union(1, 2)
        d.rollback(m2)
        self.assertTrue(d.connected(0, 1) and d.connected(2, 3))
        self.assertFalse(d.connected(0, 2))
        d.rollback(m1)
        self.assertTrue(d.connected(0, 1))
        self.assertFalse(d.connected(2, 3))

    def test_full_rollback_to_zero(self):
        d = DSU(4)
        d.union(0, 1); d.union(2, 3); d.union(1, 2)
        d.rollback(0)
        for i in range(4):
            self.assertEqual(d.size(i), 1)
        self.assertEqual(d.components(), 4)

    def test_invalid_marker(self):
        d = DSU(2)
        with self.assertRaises(ValueError):
            d.rollback(5)
        with self.assertRaises(ValueError):
            d.rollback(-1)

    def test_no_path_compression_find_stable(self):
        d = DSU(4)
        d.union(0, 1); d.union(1, 2)
        before = list(d.parent)
        r = d.find(2)
        self.assertEqual(r, d.find(0))
        self.assertEqual(d.parent[2], before[2])

    def test_size_after_rollback(self):
        d = DSU(3)
        d.union(0, 1)
        m = d.snapshot()
        d.union(0, 2)
        self.assertEqual(d.size(0), 3)
        d.rollback(m)
        self.assertEqual(d.size(0), 2)
        self.assertEqual(d.size(2), 1)

    def test_components(self):
        d = DSU(4)
        self.assertEqual(d.components(), 4)
        d.union(0, 1); d.union(2, 3)
        self.assertEqual(d.components(), 2)
        d.union(1, 2)
        self.assertEqual(d.components(), 1)

    def test_negative_n(self):
        with self.assertRaises(ValueError):
            DSU(-1)


    def test_shape_union_by_size(self):
        d = DSU(8)
        d.union(0,1); d.union(2,3); d.union(4,5); d.union(6,7)
        d.union(1,2); d.union(4,0)
        # union by size: attach smaller root under larger -> parent[4]==0, depth 2
        self.assertEqual(d.parent[4], 0)
        # deepest depth
        maxdepth = 0
        for i in range(8):
            depth=0; cur=i
            while d.parent[cur]!=cur:
                cur=d.parent[cur]; depth+=1
            maxdepth=max(maxdepth,depth)
        self.assertEqual(maxdepth, 2)

    def test_rollback_marker_type(self):
        d = DSU(2)
        for bad in (True, False, 1.5, "x"):
            with self.assertRaises((TypeError, ValueError)):
                d.rollback(bad)


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