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):
        # Q94 bullet 2: the node asked about must sit TWO levels down, else a
        # compressing find rewrites nothing and the cell is vacuous.
        # union-by-size: union(0,1)==(1<-0), union(2,3)==(3<-2), union(1,2)==(2<-0)
        d = DSU(4)
        d.union(0, 1); d.union(2, 3); d.union(1, 2)
        self.assertEqual(d.parent, [0, 0, 0, 2])
        self.assertEqual(d.parent[3], 2)          # precondition: depth 2
        before = list(d.parent)
        self.assertEqual(d.find(3), 0)            # root via parent[3] -> 2 -> 0
        self.assertEqual(d.parent, before)        # a compressing find writes parent[3]=0
        self.assertEqual(d.size(3), 4)

    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)

    def _fingerprint(self, d):
        return (list(d.parent), list(d.sz), list(d.history), d.components())

    def test_refuses_out_of_range_index(self):
        # Q94 bullet 3: every entry point that takes an index refuses a range
        # violation, and a refused call mutates nothing.
        d = DSU(3)
        d.union(0, 1)
        before = self._fingerprint(d)
        for bad in (-1, 3, 10):
            for name in ("find", "size"):
                with self.assertRaises(ValueError, msg="%s(%r)" % (name, bad)):
                    getattr(d, name)(bad)
                self.assertEqual(self._fingerprint(d), before)
            for args in ((bad, 0), (0, bad)):
                with self.assertRaises(ValueError, msg="union%r" % (args,)):
                    d.union(*args)
                self.assertEqual(self._fingerprint(d), before)
                with self.assertRaises(ValueError, msg="connected%r" % (args,)):
                    d.connected(*args)
                self.assertEqual(self._fingerprint(d), before)

    def test_refuses_non_int_index(self):
        # Q94 bullet 4: a bool is an int subclass in Python; True must not alias
        # index 1 and False must not merge index 0.
        d = DSU(3)
        before = self._fingerprint(d)
        for bad in (True, False):
            with self.assertRaises(TypeError, msg="find(%r)" % (bad,)):
                d.find(bad)
            with self.assertRaises(TypeError, msg="size(%r)" % (bad,)):
                d.size(bad)
            with self.assertRaises(TypeError, msg="union(%r, 1)" % (bad,)):
                d.union(bad, 1)
            with self.assertRaises(TypeError, msg="connected(%r, 0)" % (bad,)):
                d.connected(bad, 0)
            self.assertEqual(self._fingerprint(d), before)
        with self.assertRaises(TypeError):
            d.find(1.5)
        with self.assertRaises(TypeError):
            d.find(None)


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