import gradio as gr
import matplotlib.pyplot as plt
import numpy as np
from io import BytesIO
import networkx as nx
from collections import deque
import traceback

class DSAVisualizer:
    def __init__(self):
        self.reset_state()
    
    def reset_state(self):
        self.current_data = []
        self.steps = []
        self.current_step = 0
        self.pseudocode = []
        self.active_lines = []
        self.graph_dict = {}
        self.error_message = ""
        
    def generate_data(self, data_size=10, data_type="Random"):
        try:
            if data_size < 5:
                data_size = 5
            elif data_size > 50:
                data_size = 50
                
            if data_type == "Random":
                self.current_data = np.random.randint(1, 100, data_size)
            elif data_type == "Ascending":
                self.current_data = np.arange(1, data_size + 1)
            elif data_type == "Descending":
                self.current_data = np.arange(data_size, 0, -1)
            elif data_type == "Nearly Sorted":
                self.current_data = np.arange(1, data_size + 1)
                for _ in range(max(1, data_size // 10)):
                    i, j = np.random.randint(0, data_size, 2)
                    self.current_data[i], self.current_data[j] = self.current_data[j], self.current_data[i]
            return list(self.current_data)
        except Exception as e:
            self.error_message = f"Data generation error: {str(e)}"
            return []

    def bubble_sort(self, arr):
        try:
            arr = [int(x) for x in arr]  # Ensure integers
            n = len(arr)
            steps = [arr.copy()]
            pseudocode = [
                "procedure bubbleSort(A : list)",
                "    n = length(A)",
                "    repeat",
                "        swapped = false",
                "        for i from 1 to n-1:",
                "            if A[i-1] > A[i]:",
                "                swap(A[i-1], A[i])",
                "                swapped = true",
                "    until not swapped"
            ]
            active_line = [0]
            
            swapped = True
            while swapped:
                swapped = False
                for j in range(1, n):
                    active_line.append(5)
                    if steps[-1][j-1] > steps[-1][j]:
                        active_line.append(6)
                        new_step = steps[-1].copy()
                        new_step[j-1], new_step[j] = new_step[j], new_step[j-1]
                        steps.append(new_step)
                        swapped = True
                        active_line.append(7)
                    active_line.append(4)
                active_line.append(3)
            
            return steps, pseudocode, active_line
        except Exception as e:
            self.error_message = f"Bubble sort error: {str(e)}"
            return [], [], []

    def insertion_sort(self, arr):
        try:
            arr = [int(x) for x in arr]  # Ensure integers
            steps = [arr.copy()]
            pseudocode = [
                "procedure insertionSort(A : list)",
                "    for j from 1 to length(A)-1:",
                "        key = A[j]",
                "        i = j-1",
                "        while i >= 0 and A[i] > key:",
                "            A[i+1] = A[i]",
                "            i = i-1",
                "        A[i+1] = key"
            ]
            active_line = [0]

            arr = arr.copy()
            n = len(arr)
            for j in range(1, n):
                active_line.append(1)
                key = arr[j]
                active_line.append(2)
                i = j-1
                active_line.append(3)
                while i >= 0 and arr[i] > key:
                    active_line.append(5)
                    arr[i+1] = arr[i]
                    steps.append(arr.copy())
                    active_line.append(6)
                    i = i-1
                    active_line.append(3)
                active_line.append(7)
                arr[i+1] = key
                steps.append(arr.copy())
            
            return steps, pseudocode, active_line
        except Exception as e:
            self.error_message = f"Insertion sort error: {str(e)}"
            return [], [], []

    def dfs(self, graph_dict, start):
        try:
            graph = self.parse_graph(graph_dict)
            if not graph:
                return [], [], []
                
            visited = set()
            stack = [start]
            steps = []
            pseudocode = [
                "procedure DFS(G, start):",
                "    visited = set()",
                "    stack = [start]",
                "    while stack not empty:",
                "        vertex = stack.pop()",
                "        if vertex not in visited:",
                "            visited.add(vertex)",
                "            for neighbor in G[vertex]:",
                "                if neighbor not in visited:",
                "                    stack.push(neighbor)"
            ]
            active_line = [0]
            
            steps.append({"visited": set(), "current": None, "stack": stack.copy(), "graph": graph})
            active_line.append(1)
            
            while stack:
                active_line.append(3)
                vertex = stack.pop()
                active_line.append(4)
                if vertex not in visited:
                    active_line.append(5)
                    visited.add(vertex)
                    active_line.append(6)
                    for neighbor in graph.get(vertex, []):
                        active_line.append(7)
                        if neighbor not in visited:
                            stack.append(neighbor)
                    steps.append({"visited": visited.copy(), 
                                  "current": vertex, 
                                  "stack": stack.copy(),
                                  "graph": graph})
                active_line.append(3)
            
            return steps, pseudocode, active_line
        except Exception as e:
            self.error_message = f"DFS error: {str(e)}"
            return [], [], []

    def bfs(self, graph_dict, start):
        try:
            graph = self.parse_graph(graph_dict)
            if not graph:
                return [], [], []
                
            visited = set()
            queue = deque([start])
            steps = []
            pseudocode = [
                "procedure BFS(G, start):",
                "    visited = set()",
                "    queue = deque([start])",
                "    while queue not empty:",
                "        vertex = queue.popleft()",
                "        if vertex not in visited:",
                "            visited.add(vertex)",
                "            for neighbor in G[vertex]:",
                "                if neighbor not in visited:",
                "                    queue.append(neighbor)"
            ]
            active_line = [0]
            
            steps.append({"visited": set(), "current": None, "queue": list(queue), "graph": graph})
            active_line.append(1)
            
            while queue:
                active_line.append(3)
                vertex = queue.popleft()
                active_line.append(4)
                if vertex not in visited:
                    active_line.append(5)
                    visited.add(vertex)
                    active_line.append(6)
                    for neighbor in graph.get(vertex, []):
                        active_line.append(7)
                        if neighbor not in visited:
                            queue.append(neighbor)
                    steps.append({"visited": visited.copy(), 
                                  "current": vertex, 
                                  "queue": list(queue),
                                  "graph": graph})
                active_line.append(3)
            
            return steps, pseudocode, active_line
        except Exception as e:
            self.error_message = f"BFS error: {str(e)}"
            return [], [], []

    def visualize_sorting(self, algorithm, data):
        try:
            # Convert data to list of integers
            if isinstance(data, str):
                data = [int(np.int64(x.strip())) for x in data.strip('[]').split(',') if x.strip()]
            elif not isinstance(data, list):
                data = list(data)
                
            if algorithm == "Bubble Sort":
                self.steps, self.pseudocode, self.active_lines = self.bubble_sort(data)
            elif algorithm == "Insertion Sort":
                self.steps, self.pseudocode, self.active_lines = self.insertion_sort(data)
            else:
                self.error_message = f"Algorithm '{algorithm}' not implemented yet"
                return None, self.create_pseudocode(0), self.error_message
                
            self.current_step = 0
            return self.create_plot(), self.create_pseudocode(0), ""
        except Exception as e:
            self.error_message = f"Visualization error: {str(e)}"
            return None, "", self.error_message

    def visualize_graph(self, algorithm, graph_input, start):
        try:
            self.steps, self.pseudocode, self.active_lines = (
                self.dfs(graph_input, start) if algorithm == "DFS" 
                else self.bfs(graph_input, start)
            )
            self.current_step = 0
            return self.create_graph_plot(), self.create_pseudocode(0), ""
        except Exception as e:
            self.error_message = f"Graph visualization error: {str(e)}"
            return None, "", self.error_message

    def create_plot(self):
        try:
            fig, ax = plt.subplots(figsize=(10, 6))
            if self.steps and self.current_step < len(self.steps):
                current_data = self.steps[self.current_step]
                colors = ['#1f77b4' for _ in current_data]
                
                # Highlight recently swapped elements
                if self.current_step > 0 and self.current_step < len(self.steps):
                    prev = self.steps[self.current_step-1]
                    for i in range(len(current_data)):
                        if i < len(prev) and current_data[i] != prev[i]:
                            colors[i] = '#ff7f0e'
                
                ax.bar(range(len(current_data)), current_data, color=colors)
                ax.set_title(f'Step {self.current_step}/{len(self.steps)-1}')
                ax.set_xlabel('Index')
                ax.set_ylabel('Value')
                ax.grid(axis='y', linestyle='--', alpha=0.7)
            else:
                ax.text(0.5, 0.5, "No data to visualize", 
                        ha='center', va='center', fontsize=16)
            
            buf = BytesIO()
            plt.savefig(buf, format='png', dpi=100, bbox_inches='tight')
            plt.close(fig)
            return buf.getvalue()
        except Exception as e:
            self.error_message = f"Plot creation error: {str(e)}"
            return None

    def create_graph_plot(self):
        try:
            if not self.steps or self.current_step >= len(self.steps):
                fig, ax = plt.subplots(figsize=(10, 8))
                ax.text(0.5, 0.5, "No graph data to visualize", 
                        ha='center', va='center', fontsize=16)
                buf = BytesIO()
                plt.savefig(buf, format='png')
                plt.close(fig)
                return buf.getvalue()
                
            state = self.steps[self.current_step]
            graph_dict = state["graph"]
            
            G = nx.Graph()
            for node, neighbors in graph_dict.items():
                for neighbor in neighbors:
                    if neighbor in graph_dict:  # Ensure neighbor exists
                        G.add_edge(node, neighbor)
            
            # Add isolated nodes
            for node in graph_dict:
                if node not in G:
                    G.add_node(node)
            
            pos = nx.spring_layout(G, seed=42)
            fig, ax = plt.subplots(figsize=(10, 8))
            
            node_colors = []
            for node in G.nodes():
                if node == state.get('current'):
                    node_colors.append('#d62728')  # Current node
                elif node in state.get('visited', set()):
                    node_colors.append('#2ca02c')  # Visited nodes
                elif node in state.get('stack', []) or node in state.get('queue', []):
                    node_colors.append('#ff7f0e')  # Nodes in stack/queue
                else:
                    node_colors.append('#1f77b4')  # Unexplored nodes
            
            nx.draw_networkx_nodes(G, pos, node_size=800, 
                                   node_color=node_colors, alpha=0.9, ax=ax)
            nx.draw_networkx_edges(G, pos, width=1.5, alpha=0.5, ax=ax)
            nx.draw_networkx_labels(G, pos, font_size=12, 
                                   font_weight='bold', font_color='white', ax=ax)
            
            algorithm = "DFS" if 'stack' in state else "BFS"
            title = f"{algorithm} Traversal - Step {self.current_step}/{len(self.steps)-1}\n"
            title += f"Current: {state.get('current', 'None')} | "
            
            if 'stack' in state:
                title += f"Stack: {state['stack']}"
            elif 'queue' in state:
                title += f"Queue: {state['queue']}"
                
            ax.set_title(title, fontsize=14)
            ax.set_axis_off()
            
            buf = BytesIO()
            plt.savefig(buf, format='png', dpi=100, bbox_inches='tight')
            plt.close(fig)
            return buf.getvalue()
        except Exception as e:
            self.error_message = f"Graph plot error: {str(e)}"
            return None

    def create_pseudocode(self, step):
        try:
            if not self.pseudocode or not self.active_lines:
                return ""
            
            # Find the active line for this step
            if step < len(self.active_lines):
                active_line = self.active_lines[step]
            else:
                active_line = self.active_lines[-1] if self.active_lines else 0
            
            html = "<div style='font-family: monospace; background: #2d2d2d; color: #f8f8f2; padding: 15px; border-radius: 8px; line-height: 1.5;'>"
            for i, line in enumerate(self.pseudocode):
                if i == active_line:
                    html += f"<div style='background: #44475a; padding: 8px; border-left: 3px solid #bd93f9;'><b>{line}</b></div>"
                else:
                    html += f"<div style='padding: 8px;'>{line}</div>"
            html += "</div>"
            return html
        except Exception as e:
            self.error_message = f"Pseudocode error: {str(e)}"
            return ""

    def next_step(self):
        if self.current_step < len(self.steps) - 1:
            self.current_step += 1
        return self.update_display()
    
    def prev_step(self):
        if self.current_step > 0:
            self.current_step -= 1
        return self.update_display()
    
    def update_display(self):
        try:
            if not self.steps:
                return None, "", "No steps available"
            
            if isinstance(self.steps[0], list):  # Sorting visualization
                plot = self.create_plot()
            else:  # Graph visualization
                plot = self.create_graph_plot()
            
            return plot, self.create_pseudocode(self.current_step), ""
        except Exception as e:
            return None, "", f"Update error: {str(e)}"

    def parse_graph(self, graph_str):
        try:
            graph = {}
            for line in graph_str.strip().split('\n'):
                if ':' in line:
                    node, neighbors = line.split(':', 1)
                    node = node.strip()
                    if node not in graph:
                        graph[node] = []
                    neighbors = [n.strip() for n in neighbors.split(',') if n.strip()]
                    graph[node].extend(neighbors)
                    
                    # Add neighbors that might not be defined yet
                    for neighbor in neighbors:
                        if neighbor not in graph:
                            graph[neighbor] = []
            return graph
        except Exception as e:
            self.error_message = f"Graph parse error: {str(e)}"
            return {}

# Initialize visualizer
visualizer = DSAVisualizer()

# CSS for styling
custom_css = """
:root {
    --primary: #6e40c9;
    --secondary: #3d5afe;
    --accent: #ff4081;
    --dark: #2d2d2d;
    --light: #f8f8f8;
    --success: #4CAF50;
    --error: #f44336;
}

body {
    background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%);
    color: var(--light);
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
}

h1, h2, h3 {
    background: linear-gradient(90deg, var(--primary), var(--secondary));
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    text-shadow: 0 2px 4px rgba(0,0,0,0.1);
}

.tab {
    background: rgba(45, 45, 45, 0.8) !important;
    backdrop-filter: blur(10px);
    border-radius: 12px;
    padding: 20px;
    box-shadow: 0 8px 32px rgba(0,0,0,0.2);
    border: 1px solid rgba(255,255,255,0.1);
}

.btn-primary {
    background: linear-gradient(135deg, var(--primary), var(--secondary)) !important;
    border: none !important;
    border-radius: 8px !important;
    font-weight: 600 !important;
    text-transform: uppercase !important;
    letter-spacing: 0.5px !important;
    transition: all 0.3s ease !important;
    color: white !important;
}

.btn-primary:hover {
    transform: translateY(-2px);
    box-shadow: 0 6px 12px rgba(110, 64, 201, 0.3) !important;
}

.control-panel {
    background: rgba(45, 45, 45, 0.6) !important;
    border-radius: 12px;
    padding: 20px;
    border: 1px solid rgba(255,255,255,0.1);
}

.visualization-container {
    background: rgba(45, 45, 45, 0.6);
    border-radius: 12px;
    padding: 20px;
    border: 1px solid rgba(255,255,255,0.1);
    min-height: 500px;
}

.error-box {
    background: var(--error) !important;
    color: white !important;
    padding: 10px;
    border-radius: 8px;
    margin-top: 10px;
    font-weight: bold;
}

.success-box {
    background: var(--success) !important;
    color: white !important;
    padding: 10px;
    border-radius: 8px;
    margin-top: 10px;
    font-weight: bold;
}

footer {
    text-align: center;
    margin-top: 20px;
    color: rgba(255,255,255,0.6);
    font-size: 0.9em;
}
"""

# Gradio Interface
with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple"), 
              css=custom_css, 
              title="AlgoViz Pro") as demo:
    
    # Header
    with gr.Row():
        gr.Markdown("""
        <div style="text-align: center; width: 100%;">
            <h1 style="font-size: 2.5rem; margin-bottom: 0.5rem;">🌟 AlgoViz Pro</h1>
            <p style="font-size: 1.2rem; color: #a0a0c0; max-width: 800px; margin: 0 auto;">
                Interactive Data Structures & Algorithms Visualization Platform
            </p>
        </div>
        """)
    
    # Error display
    error_output = gr.HTML(visible=False, elem_classes=["error-box"])
    success_output = gr.HTML(visible=False, elem_classes=["success-box"])
    
    # Main Content
    with gr.Tab("🧮 Sorting Algorithms"):
        with gr.Row():
            # Control Panel
            with gr.Column(scale=1, elem_classes="control-panel"):
                data_type = gr.Radio(
                    choices=["Random", "Ascending", "Descending", "Nearly Sorted"],
                    value="Random", label="Data Distribution"
                )
                data_size = gr.Slider(5, 50, value=10, step=1, label="Data Size")
                gen_btn = gr.Button("♻️ Generate Data", variant="primary")
                algorithm = gr.Dropdown(
                    ["Bubble Sort", "Insertion Sort"],
                    value="Bubble Sort", label="Algorithm"
                )
                sort_btn = gr.Button("▶️ Run Algorithm", variant="primary")
                
                data_output = gr.Textbox(label="Generated Data", interactive=True)
                
                gen_btn.click(
                    visualizer.generate_data, 
                    inputs=[data_size, data_type], 
                    outputs=[data_output]
                )
                
            # Visualization Area
            with gr.Column(scale=2, elem_classes="visualization-container"):
                plot_output = gr.Image(label="Visualization", height=400)
                pseudocode = gr.HTML(label="Pseudocode")
                
                with gr.Row():
                    prev_btn = gr.Button("⏪ Previous Step")
                    next_btn = gr.Button("⏩ Next Step")
                    step_counter = gr.Number(value=0, label="Current Step", interactive=False)
        
        sort_btn.click(
            visualizer.visualize_sorting,
            inputs=[algorithm, data_output],
            outputs=[plot_output, pseudocode, error_output]
        )
    
    with gr.Tab("📊 Graph Algorithms"):
        with gr.Row():
            # Control Panel
            with gr.Column(scale=1, elem_classes="control-panel"):
                gr.Markdown("### Graph Definition")
                graph_input = gr.Textbox(
                    value="A: B,C\nB: A,D\nC: A,E\nD: B\nE: C",
                    lines=7,
                    label="Adjacency List",
                    placeholder="Enter graph as:\nnode: neighbor1,neighbor2\n...",
                )
                with gr.Row():
                    start_node = gr.Textbox(value="A", label="Start Node")
                    algorithm = gr.Dropdown(
                        ["DFS", "BFS"], 
                        value="DFS", label="Algorithm"
                    )
                graph_btn = gr.Button("▶️ Run Algorithm", variant="primary")
                
            # Visualization Area
            with gr.Column(scale=2, elem_classes="visualization-container"):
                graph_plot = gr.Image(label="Graph Visualization", height=400)
                graph_pseudocode = gr.HTML(label="Pseudocode")
                
                with gr.Row():
                    graph_prev = gr.Button("⏪ Previous Step")
                    graph_next = gr.Button("⏩ Next Step")
                    graph_step = gr.Number(value=0, label="Current Step", interactive=False)
        
        graph_btn.click(
            visualizer.visualize_graph,
            inputs=[algorithm, graph_input, start_node],
            outputs=[graph_plot, graph_pseudocode, error_output]
        )
    
    with gr.Tab("📚 Complexity Analysis"):
        gr.Markdown("""
        ## Algorithm Complexity Cheat Sheet
        
        | Algorithm       | Time (Best)   | Time (Avg)    | Time (Worst)  | Space         | Use Cases                     |
        |-----------------|---------------|---------------|---------------|---------------|-------------------------------|
        | **Bubble Sort** | O(n)          | O(n²)         | O(n²)         | O(1)          | Educational, small datasets  |
        | **Insertion Sort** | O(n)       | O(n²)         | O(n²)         | O(1)          | Small datasets, nearly sorted|
        | **Merge Sort**  | O(n log n)    | O(n log n)    | O(n log n)    | O(n)          | General purpose, stable      |
        | **Quick Sort**  | O(n log n)    | O(n log n)    | O(n²)         | O(log n)      | Large datasets, in-place     |
        | **DFS**         | O(V+E)        | O(V+E)        | O(V+E)        | O(V)          | Pathfinding, cycle detection |
        | **BFS**         | O(V+E)        | O(V+E)        | O(V+E)        | O(V)          | Shortest path, level order   |
        
        ### Key Insights:
        - **Sorting Algorithms**: 
            - Use **QuickSort** for average-case performance
            - Use **MergeSort** for guaranteed O(n log n) performance
            - **BubbleSort** and **InsertionSort** are efficient for small datasets
        
        - **Graph Algorithms**:
            - **DFS** is better for pathfinding in deep graphs
            - **BFS** is better for shortest path in unweighted graphs
        """)
    
    with gr.Tab("💡 About"):
        gr.Markdown("""
        ## About AlgoViz Pro
        
        AlgoViz Pro is an interactive visualization tool for understanding Data Structures and Algorithms. 
        It provides step-by-step visualizations with pseudocode execution highlighting to help students 
        and developers understand how algorithms work.
        
        **Features**:
        - Real-time algorithm visualization
        - Step-by-step execution control
        - Multiple algorithm implementations
        - Pseudocode with execution highlighting
        - Customizable input data
        - Comprehensive error handling
        - Complexity analysis reference
        
        **Algorithms Implemented**:
        - Sorting: Bubble Sort, Insertion Sort
        - Graph Traversal: DFS, BFS
        
        **Future Enhancements**:
        - Add more algorithms (Merge Sort, Dijkstra, etc.)
        - Speed control for animations
        - Comparison mode for algorithms
        - Export visualizations as GIF
        
        Developed with ❤️ using Python, Gradio, and Matplotlib
        """)
    
    # Navigation events
    next_btn.click(
        visualizer.next_step, 
        outputs=[plot_output, pseudocode, error_output]
    )
    prev_btn.click(
        visualizer.prev_step, 
        outputs=[plot_output, pseudocode, error_output]
    )
    graph_next.click(
        visualizer.next_step, 
        outputs=[graph_plot, graph_pseudocode, error_output]
    )
    graph_prev.click(
        visualizer.prev_step, 
        outputs=[graph_plot, graph_pseudocode, error_output]
    )
    
    # Footer
    gr.Markdown("""
    <footer>
        <p>AlgoViz Pro v1.1 | Robust & Error-Resistant | For educational purposes</p>
    </footer>
    """)
    
    # Error handling display
    def show_error(error_msg):
        if error_msg:
            return gr.update(value=f"⚠️ ERROR: {error_msg}", visible=True)
        return gr.update(visible=False)
    
    error_output.change(
        show_error, 
        inputs=[error_output], 
        outputs=[error_output]
    )

if __name__ == "__main__":
    demo.launch()