from collections import deque

import matplotlib.pyplot as plt
import networkx as nx

# ==========================================
# 1. THE SEARCH ALGORITHMS (Add new ones here)
# ==========================================


def bfs_algorithm(graph, start_node):
    """
    Breadth-First Search
    """
    visited = set()
    queue = deque([start_node])

    while queue:
        current = queue.popleft()

        if current not in visited:
            visited.add(current)

            for neighbor in graph.neighbors(current):
                if neighbor not in visited:
                    queue.append(neighbor)

            # Send current engine state back to the visualizer
            yield current, visited


def recursive_dfs_algorithm(graph, start_node):
    """
    Depth-First Search
    """
    visited = set()

    # 1. Define the internal recursive function
    def dfs_helper(node):
        if node not in visited:
            visited.add(node)

            # Send current frame state back up the chain to the visualizer
            yield node, visited

            for neighbor in graph.neighbors(node):
                if neighbor not in visited:
                    # 2. 'yield from' intercepts values from deep recursive steps
                    yield from dfs_helper(neighbor)

    # 3. Start the recursion and pass the generator out
    return dfs_helper(start_node)


# ==========================================
# 2. THE VISUALIZER ENGINE
# ==========================================

class GraphVisualizer:
    def __init__(self, graph):
        self.graph = graph
        # Precompute layouts so nodes stay completely still
        self.pos = nx.spring_layout(graph, seed=42)

    def run_animation(self, algorithm_generator, title="Search Animation", delay=1.5):
        """
        Consumes states from the algorithm generator and updates the frame.
        """
        plt.ion()  # Turn on interactive plotting mode
        fig, ax = plt.subplots(figsize=(6, 5))

        # Pull each state frame out of the search algorithm sequentially
        for current_node, visited_nodes in algorithm_generator:

            # Build node colors based on current frame data
            color_map = []
            for node in self.graph.nodes():
                if node == current_node:
                    color_map.append('orange')  # Current frontier node
                elif node in visited_nodes:
                    color_map.append('lightgreen')  # Already processed
                else:
                    color_map.append('skyblue')  # Unvisited

            # Draw frame
            ax.clear()
            nx.draw(
                self.graph, self.pos, with_labels=True,
                node_color=color_map, node_size=800,
                font_weight='bold', ax=ax
            )
            ax.set_title(f"{title}\nActive Node: {current_node}")

            plt.draw()
            plt.pause(delay)

        plt.ioff()
        plt.show()


# ==========================================
# 3. RUNNING THE CODE
# ==========================================
if __name__ == "__main__":
    # Define a test graph
    G = nx.Graph()
    edges = [('A', 'B'), ('A', 'C'), ('B', 'D'), ('B', 'E'), ('C', 'F'), ('E', 'F')]
    G.add_edges_from(edges)

    # Initialize visualizer engine
    viz = GraphVisualizer(G)

    # Easily swap algorithms by passing different generator objects
    print("Running Breadth-First Search...")
    bfs_gen = bfs_algorithm(G, 'A')
    viz.run_animation(bfs_gen, title="Breadth-First Search (BFS)", delay=1.0)

    print("Running Depth-First Search...")
    dfs_gen = recursive_dfs_algorithm(G, 'A')
    viz.run_animation(dfs_gen, title="Depth-First Search (DFS)", delay=1.0)