import math
import random
import unittest
from example import dijkstra


class WeightedTests(unittest.TestCase):
    def test_stale_zero_parallel_unreachable(self):
        graph = {0: [(1, 9), (1, 8), (2, 0)], 2: [(1, 2)], 1: [], 3: []}
        self.assertEqual(dijkstra(graph, 0), {0: 0, 1: 2, 2: 0, 3: math.inf})

    def test_invalid_weights(self):
        for weight in (-1, math.inf, math.nan):
            with self.assertRaises(ValueError):
                dijkstra({0: [(1, weight)]}, 0)

    def test_against_floyd(self):
        rng = random.Random(14)
        for _ in range(30):
            n = 6
            graph = {i: [] for i in range(n)}
            dist = [[0 if i == j else math.inf for j in range(n)] for i in range(n)]
            for i in range(n):
                for j in range(n):
                    if i != j and rng.random() < 0.25:
                        weight = rng.randrange(10)
                        graph[i].append((j, weight))
                        dist[i][j] = weight
            for k in range(n):
                for i in range(n):
                    for j in range(n):
                        dist[i][j] = min(dist[i][j], dist[i][k] + dist[k][j])
            for source in range(n):
                self.assertEqual(dijkstra(graph, source), dict(enumerate(dist[source])))


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