import unittest
from floydwarshall import floyd


class TestFloyd(unittest.TestCase):

    def test_floyd_1(self):
        inf = float('inf')
        g = {0: {0: 0,   1: 1,   2: 4},
             1: {0: inf, 1: 0,  2: 2},
             2: {0: inf, 1: inf, 2: 0}}

        expected = {0: {0: 0, 1: 1, 2: 3},
                    1: {0: inf, 1: 0, 2: 2},
                    2: {0: inf, 1: inf, 2: 0}}

        paths = floyd(g)
        self.assertDictEqual(paths, expected)
        print paths

    def test_floyd_2(self):

        def _adj(g):   # Convert a directed graph to an adjacency matrix.
            vertices = g.keys()
            return {v1: {v2: 0 if v1 == v2 else g[v1].get(v2, float('inf'))
                         for v2 in vertices}
                    for v1 in vertices}

        g = {1: {2: 3, 3: 8, 5: -4},
             2: {4: 1, 5: 7},
             3: {2: 4},
             4: {1: 2, 3: -5},
             5: {4: 6}}

        expected = {1: {1: 0, 2: 1, 3: -3, 4: 2, 5: -4},
                    2: {1: 3, 2: 0, 3: -4, 4: 1, 5: -1},
                    3: {1: 7, 2: 4, 3: 0, 4: 5, 5: 3},
                    4: {1: 2, 2: -1, 3: -5, 4: 0, 5: -2},
                    5: {1: 8, 2: 5, 3: 1, 4: 6, 5: 0}}

        paths = floyd(_adj(g))
        self.assertDictEqual(paths, expected)
        print paths

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