#!/usr/bin/env python
import pytest
import networkx as nx


class TestFloyd:
    @classmethod
    def setup_class(cls):
        pass

    def test_floyd_warshall_predecessor_and_distance(self):
        XG = nx.DiGraph()
        XG.add_weighted_edges_from([('s', 'u', 10), ('s', 'x', 5),
                                    ('u', 'v', 1), ('u', 'x', 2),
                                    ('v', 'y', 1), ('x', 'u', 3),
                                    ('x', 'v', 5), ('x', 'y', 2),
                                    ('y', 's', 7), ('y', 'v', 6)])
        path, dist = nx.floyd_warshall_predecessor_and_distance(XG)
        assert dist['s']['v'] == 9
        assert path['s']['v'] == 'u'
        assert (dist ==
                     {'y': {'y': 0, 'x': 12, 's': 7, 'u': 15, 'v': 6},
                      'x': {'y': 2, 'x': 0, 's': 9, 'u': 3, 'v': 4},
                      's': {'y': 7, 'x': 5, 's': 0, 'u': 8, 'v': 9},
                      'u': {'y': 2, 'x': 2, 's': 9, 'u': 0, 'v': 1},
                      'v': {'y': 1, 'x': 13, 's': 8, 'u': 16, 'v': 0}})

        GG = XG.to_undirected()
        # make sure we get lower weight
        # to_undirected might choose either edge with weight 2 or weight 3
        GG['u']['x']['weight'] = 2
        path, dist = nx.floyd_warshall_predecessor_and_distance(GG)
        assert dist['s']['v'] == 8
        # skip this test, could be alternate path s-u-v
#        assert_equal(path['s']['v'],'y')

        G = nx.DiGraph()  # no weights
        G.add_edges_from([('s', 'u'), ('s', 'x'),
                          ('u', 'v'), ('u', 'x'),
                          ('v', 'y'), ('x', 'u'),
                          ('x', 'v'), ('x', 'y'),
                          ('y', 's'), ('y', 'v')])
        path, dist = nx.floyd_warshall_predecessor_and_distance(G)
        assert dist['s']['v'] == 2
        # skip this test, could be alternate path s-u-v
        # assert_equal(path['s']['v'],'x')

        # alternate interface
        dist = nx.floyd_warshall(G)
        assert dist['s']['v'] == 2

        # floyd_warshall_predecessor_and_distance returns
        # dicts-of-defautdicts
        # make sure we don't get empty dictionary
        XG = nx.DiGraph()
        XG.add_weighted_edges_from([('v', 'x', 5.0), ('y', 'x', 5.0),
                                    ('v', 'y', 6.0), ('x', 'u', 2.0)])
        path, dist = nx.floyd_warshall_predecessor_and_distance(XG)
        inf = float("inf")
        assert (dist ==
                     {'v': {'v': 0, 'x': 5.0, 'y': 6.0, 'u': 7.0},
                      'x': {'x': 0, 'u': 2.0, 'v': inf, 'y': inf},
                      'y': {'y': 0, 'x': 5.0, 'v': inf, 'u': 7.0},
                      'u': {'u': 0, 'v': inf, 'x': inf, 'y': inf}})
        assert (path ==
                     {'v': {'x': 'v', 'y': 'v', 'u': 'x'},
                      'x': {'u': 'x'},
                      'y': {'x': 'y', 'u': 'x'}})

    def test_reconstruct_path(self):
        with pytest.raises(KeyError):
            XG = nx.DiGraph()
            XG.add_weighted_edges_from([('s', 'u', 10), ('s', 'x', 5),
                                        ('u', 'v', 1), ('u', 'x', 2),
                                        ('v', 'y', 1), ('x', 'u', 3),
                                        ('x', 'v', 5), ('x', 'y', 2),
                                        ('y', 's', 7), ('y', 'v', 6)])
            predecessors, _ = nx.floyd_warshall_predecessor_and_distance(XG)

            path = nx.reconstruct_path('s', 'v', predecessors)
            assert path == ['s', 'x', 'u', 'v']

            path = nx.reconstruct_path('s', 's', predecessors)
            assert path == []

            # this part raises the keyError
            nx.reconstruct_path('1', '2', predecessors)

    def test_cycle(self):
        path, dist = nx.floyd_warshall_predecessor_and_distance(
            nx.cycle_graph(7))
        assert dist[0][3] == 3
        assert path[0][3] == 2
        assert dist[0][4] == 3

    def test_weighted(self):
        XG3 = nx.Graph()
        XG3.add_weighted_edges_from([[0, 1, 2], [1, 2, 12], [2, 3, 1],
                                     [3, 4, 5], [4, 5, 1], [5, 0, 10]])
        path, dist = nx.floyd_warshall_predecessor_and_distance(XG3)
        assert dist[0][3] == 15
        assert path[0][3] == 2

    def test_weighted2(self):
        XG4 = nx.Graph()
        XG4.add_weighted_edges_from([[0, 1, 2], [1, 2, 2], [2, 3, 1],
                                     [3, 4, 1], [4, 5, 1], [5, 6, 1],
                                     [6, 7, 1], [7, 0, 1]])
        path, dist = nx.floyd_warshall_predecessor_and_distance(XG4)
        assert dist[0][2] == 4
        assert path[0][2] == 1

    def test_weight_parameter(self):
        XG4 = nx.Graph()
        XG4.add_edges_from([(0, 1, {'heavy': 2}), (1, 2, {'heavy': 2}),
                            (2, 3, {'heavy': 1}), (3, 4, {'heavy': 1}),
                            (4, 5, {'heavy': 1}), (5, 6, {'heavy': 1}),
                            (6, 7, {'heavy': 1}), (7, 0, {'heavy': 1})])
        path, dist = nx.floyd_warshall_predecessor_and_distance(XG4,
                                                                weight='heavy')
        assert dist[0][2] == 4
        assert path[0][2] == 1

    def test_zero_distance(self):
        XG = nx.DiGraph()
        XG.add_weighted_edges_from([('s', 'u', 10), ('s', 'x', 5),
                                    ('u', 'v', 1), ('u', 'x', 2),
                                    ('v', 'y', 1), ('x', 'u', 3),
                                    ('x', 'v', 5), ('x', 'y', 2),
                                    ('y', 's', 7), ('y', 'v', 6)])
        path, dist = nx.floyd_warshall_predecessor_and_distance(XG)

        for u in XG:
            assert dist[u][u] == 0

        GG = XG.to_undirected()
        # make sure we get lower weight
        # to_undirected might choose either edge with weight 2 or weight 3
        GG['u']['x']['weight'] = 2
        path, dist = nx.floyd_warshall_predecessor_and_distance(GG)

        for u in GG:
            dist[u][u] = 0

    def test_zero_weight(self):
        G = nx.DiGraph()
        edges = [(1, 2, -2), (2, 3, -4), (1, 5, 1),
                 (5, 4, 0), (4, 3, -5), (2, 5, -7)]
        G.add_weighted_edges_from(edges)
        dist = nx.floyd_warshall(G)
        assert dist[1][3] == -14

        G = nx.MultiDiGraph()
        edges.append((2, 5, -7))
        G.add_weighted_edges_from(edges)
        dist = nx.floyd_warshall(G)
        assert dist[1][3] == -14