from dijkstra import dijkstra, reconstruct_path
import random

# To run these tests, use `pytest test_dijkstra.py`.
# Disclosure: I used Copilot to help me remember how to generate random 
# numbers in Python and how to use Pytest. I also used it to find a good
# priority-queue library in Python (I didn't know there was one built in!)
#   - Tim

class TestBasicShortestPaths:
    """Test basic shortest path functionality."""
    
    EXAMPLE_GRAPH = {
        'G': [('A', 3), ('D', 4), ('C', 3)],
        'A': [('G', 3), ('B', 1), ('D', 3)],
        'D': [('A', 3), ('G', 4), ('C', 2), ('B', 5)],
        'B': [('D', 5), ('A', 1), ('E', 2)],
        'E': [('B', 2), ('F', 2)],
        'F': [('E', 2), ('C', 6)],
        'C': [('G', 3), ('D', 2), ('F', 6)]
    }

    def test_shortest_distances_from_g(self):
        """Test that shortest distances from G are computed correctly."""
        distances, predecessors = dijkstra(self.EXAMPLE_GRAPH, 'G')
        
        assert distances['G'] == 0, f"Distance to G should be 0, got {distances['G']}"
        assert distances['A'] == 3, f"Distance to A should be 3, got {distances['A']}"
        assert distances['B'] == 4, f"Distance to B should be 4, got {distances['B']}"
        assert distances['E'] == 6, f"Distance to E should be 6, got {distances['E']}"
        assert distances['F'] == 8, f"Distance to F should be 8, got {distances['E']}"
    
    def test_path_reconstruction(self):
        """Test that paths are reconstructed correctly."""
        # Don't call Dijkstra here, we're testing reconstruct_path alone.
        predecessors = {'E': 'B', '10': '11', 'B': 'A', 'X': 'Y', 'A': 'G', 'G': None}
        
        path_to_e = reconstruct_path(predecessors, 'G', 'E')
        expected_path = ['G', 'A', 'B', 'E']
        assert path_to_e == expected_path, f"Path to E should be {expected_path}, got {path_to_e}"
    
    def test_path_to_self(self):
        """Test that path to starting node is just the node itself."""
        distances, predecessors = dijkstra(self.EXAMPLE_GRAPH, 'G')
        
        path_to_g = reconstruct_path(predecessors, 'G', 'G')
        assert path_to_g == ['G'], f"Path to self should be ['G'], got {path_to_g}"

    def test_start_from_b(self):
        """Test starting from a node other than G."""
        distances, predecessors = dijkstra(self.EXAMPLE_GRAPH, 'B')
        
        assert distances['B'] == 0, "Distance to start should be 0"
        
        path_to_g = reconstruct_path(predecessors, 'B', 'G')
        assert path_to_g is not None, "Path from B to G should exist"
        assert path_to_g[0] == 'B', "Path should start with B"
        assert path_to_g[-1] == 'G', "Path should end with G"

class TestSingleNodeGraph:
    """Test edge case of single node graph."""
    
    def test_single_node(self):
        """Test that single node graph works correctly."""
        single_graph = {'X': []}
        distances, predecessors = dijkstra(single_graph, 'X')
        
        assert distances['X'] == 0, "Distance to itself should be 0"
        
        path = reconstruct_path(predecessors, 'X', 'X')
        assert path == ['X'], f"Path to itself should be ['X'], got {path}"


class TestDisconnectedGraph:
    """Test graphs with unreachable nodes."""
    
    disconnected_graph = {
            'A': [('B', 1)],
            'B': [('A', 1)],
            'C': [('D', 1)],
            'D': [('C', 1)]
        }

    def test_unreachable_nodes(self):
        """Test that unreachable nodes have infinite distance."""
        distances, predecessors = dijkstra(self.disconnected_graph, 'A')
        
        assert distances['A'] == 0, "Distance to start should be 0"
        assert distances['B'] == 1, "Distance to B should be 1"
        assert distances['C'] == float('inf'), "Distance to C should be infinity"
        assert distances['D'] == float('inf'), "Distance to D should be infinity"
    
    def test_no_path_returns_none(self):
        """Test that reconstruct_path returns None for unreachable nodes."""
        distances, predecessors = dijkstra(self.disconnected_graph, 'A')
        
        path_to_c = reconstruct_path(predecessors, 'A', 'C')
        assert path_to_c is None, "Path to unreachable node should be None"


# Note that this test is a stepping-stone on the way to Property-Based Testing!
# We're recognizing that multiple paths might exist, and checking specific
# *properties* of the result. We're never checking for a single _concrete_ path. 
#   (We could do much better than this, though.)
class TestMultipleEqualPaths:
    """Test graphs where multiple shortest paths exist."""
    
    def test_triangle_equal_cost_paths(self):
        """Test that algorithm finds *a* valid shortest path."""
        triangle = {
            'A': [('B', 1), ('C', 2)],
            'B': [('C', 1), ('A', 1)],
            'C': [('B', 1), ('A', 2)]
        }
        distances, predecessors = dijkstra(triangle, 'A')
        
        # PROPERTY 1: The reported cost is as expected
        # Both A->B->C and A->C have the same cost (2)
        assert distances['C'] == 2, "Distance to C should be 2"
        
        path_to_c = reconstruct_path(predecessors, 'A', 'C')
        # PROPERTY 2: Check path is valid
        assert path_to_c is not None, "Path should exist"
        # PROPERTY 3: Check path endpoints are correct
        assert path_to_c[0] == 'A', "Path should start with A"
        assert path_to_c[-1] == 'C', "Path should end with C"
        
        # PROPERTY 4: the actual path-cost total is as expected
        path_cost = 0
        for i in range(len(path_to_c) - 1):
            curr_node = path_to_c[i]
            next_node = path_to_c[i + 1]
            # Find the edge cost
            for neighbor, weight in triangle[curr_node]:
                if neighbor == next_node:
                    path_cost += weight
                    break
        assert path_cost == 2, f"Path cost to C should be 2, got {path_cost}"

# In fact, let's try to do better. :-) This is incomplete, but demonstrates the idea.
class TestRandomGraphsPBT:
    """Test on a number of randomly-generated graphs."""

    # Check on this many random graphs every time the test suite is called.
    # Setting this to a low number, because the properties below are kinda trivial.
    N_GRAPHS = 100
    
    def test_random_graphs(self):
        for _ in range(self.N_GRAPHS):
            graph = generate_random_graph(
                num_nodes=random.randint(5, 20),
                min_weight=1,
                max_weight=100,
                edge_prob=random.uniform(0.0, 1.0)
            )
            start = random.randint(0, len(graph) - 1)
            distances, predecessors = dijkstra(graph, start)
            # PROPERTY: distance to origin 
            assert distances[start] == 0, f"Distance to start node {start} should be 0"
            # ... We could add more here! 
            
def generate_random_graph(num_nodes, min_weight, max_weight, edge_prob=0.3):
    graph = {i: [] for i in range(num_nodes)}
    for i in range(num_nodes):
        for j in range(i + 1, num_nodes):
            if random.random() < edge_prob:
                weight = random.randint(min_weight, max_weight)
                graph[i].append((j, weight))
                graph[j].append((i, weight))
    return graph