"""
USSU Algorithm Analyzer v4.0 - Graph Algorithms Suite
BFS, DFS, Shortest Path, MST, DAG, SCC, Topological Sort, and more.
"""
import heapq
import time
from collections import deque, defaultdict
from typing import Dict, List, Set, Tuple, Optional, Any
from utils.core import profile_algorithm, Graph, OperationCounter


class GraphAlgorithms(OperationCounter):
    """Complete graph algorithm suite with step tracking"""

    def __init__(self):
        super().__init__()
        self.traversal_steps: List[Dict] = []

    def reset(self):
        self.reset_counters()
        self.traversal_steps = []

    def _make_result(self, name: str, time_c: str, space_c: str, **kwargs) -> Dict:
        base = {
            'algorithm': name,
            'time_complexity': time_c,
            'space_complexity': space_c,
            'comparisons': self.comparisons,
            'accesses': self.accesses,
            'iterations': self.iterations,
        }
        base.update(kwargs)
        return base

    # ==================== BFS ====================
    @profile_algorithm
    def bfs(self, graph: Graph, start: int, target: Optional[int] = None) -> Dict:
        self.reset()
        visited = set()
        queue = deque([(start, [start])])
        traversal = []
        levels = {start: 0}
        parent = {start: None}

        while queue:
            self.iterations += 1
            vertex, path = queue.popleft()
            if vertex not in visited:
                visited.add(vertex)
                traversal.append(vertex)
                self.traversal_steps.append({'node': vertex, 'frontier': list(queue), 'visited': list(visited)})
                if target is not None and vertex == target:
                    return self._make_result('BFS', 'O(V + E)', 'O(V)',
                        traversal=traversal, path=path, levels=levels, parent=parent,
                        found=True, distance=len(path)-1, vertices=len(graph.vertices), edges=graph.edge_count)
                for neighbor, _ in graph.get_neighbors(vertex):
                    self.accesses += 1
                    if neighbor not in visited:
                        queue.append((neighbor, path + [neighbor]))
                        levels[neighbor] = levels[vertex] + 1
                        parent[neighbor] = vertex
        return self._make_result('BFS', 'O(V + E)', 'O(V)',
            traversal=traversal, path=[], levels=levels, parent=parent,
            found=False, vertices=len(graph.vertices), edges=graph.edge_count)

    # ==================== DFS ====================
    @profile_algorithm
    def dfs_iterative(self, graph: Graph, start: int, target: Optional[int] = None) -> Dict:
        self.reset()
        visited = set()
        stack = [(start, [start])]
        traversal = []

        while stack:
            self.iterations += 1
            vertex, path = stack.pop()
            if vertex not in visited:
                visited.add(vertex)
                traversal.append(vertex)
                self.traversal_steps.append({'node': vertex, 'frontier': list(stack), 'visited': list(visited)})
                if target is not None and vertex == target:
                    return self._make_result('DFS (Iterative)', 'O(V + E)', 'O(V)',
                        traversal=traversal, path=path, found=True, vertices=len(graph.vertices), edges=graph.edge_count)
                for neighbor, _ in reversed(graph.get_neighbors(vertex)):
                    self.accesses += 1
                    if neighbor not in visited:
                        stack.append((neighbor, path + [neighbor]))
        return self._make_result('DFS (Iterative)', 'O(V + E)', 'O(V)',
            traversal=traversal, path=[], found=False, vertices=len(graph.vertices), edges=graph.edge_count)

    @profile_algorithm
    def dfs_recursive(self, graph: Graph, start: int, target: Optional[int] = None) -> Dict:
        self.reset()
        visited = set()
        traversal = []
        path_found = []

        def dfs_helper(vertex, path):
            self.recursions += 1
            if vertex in visited:
                return False
            visited.add(vertex)
            traversal.append(vertex)
            self.traversal_steps.append({'node': vertex, 'visited': list(visited)})
            if target is not None and vertex == target:
                path_found.extend(path)
                return True
            for neighbor, _ in graph.get_neighbors(vertex):
                self.accesses += 1
                if dfs_helper(neighbor, path + [neighbor]):
                    return True
            return False

        found = dfs_helper(start, [start])
        return self._make_result('DFS (Recursive)', 'O(V + E)', 'O(V)',
            traversal=traversal, path=path_found if found else [],
            found=found, vertices=len(graph.vertices), edges=graph.edge_count)

    # ==================== DIJKSTRA ====================
    @profile_algorithm
    def dijkstra(self, graph: Graph, start: int, target: Optional[int] = None) -> Dict:
        self.reset()
        if not graph.weighted:
            return self.bfs(graph, start, target)
        distances = {v: float('inf') for v in graph.vertices}
        distances[start] = 0
        parent = {v: None for v in graph.vertices}
        visited = set()
        pq = [(0, start)]

        while pq:
            self.iterations += 1
            dist, vertex = heapq.heappop(pq)
            if vertex in visited:
                continue
            visited.add(vertex)
            if target is not None and vertex == target:
                break
            for neighbor, weight in graph.get_neighbors(vertex):
                self.accesses += 1
                self.comparisons += 1
                w = weight if weight is not None else 1
                new_dist = dist + w
                if new_dist < distances[neighbor]:
                    distances[neighbor] = new_dist
                    parent[neighbor] = vertex
                    heapq.heappush(pq, (new_dist, neighbor))

        # Reconstruct path
        path = []
        if target is not None and distances[target] != float('inf'):
            curr = target
            while curr is not None:
                path.append(curr)
                curr = parent[curr]
            path.reverse()
        return self._make_result('Dijkstra', 'O((V+E) log V)', 'O(V)',
            distances=distances, parent=parent, path=path,
            vertices=len(graph.vertices), edges=graph.edge_count)

    # ==================== BELLMAN-FORD ====================
    @profile_algorithm
    def bellman_ford(self, graph: Graph, start: int) -> Dict:
        self.reset()
        distances = {v: float('inf') for v in graph.vertices}
        distances[start] = 0
        parent = {v: None for v in graph.vertices}
        V = len(graph.vertices)

        for i in range(V - 1):
            updated = False
            for u, v, w in graph.edges:
                weight = w if w is not None else 1
                if distances[u] != float('inf') and distances[u] + weight < distances[v]:
                    distances[v] = distances[u] + weight
                    parent[v] = u
                    updated = True
                    self.comparisons += 1
            if not updated:
                break

        # Negative cycle check
        for u, v, w in graph.edges:
            weight = w if w is not None else 1
            if distances[u] != float('inf') and distances[u] + weight < distances[v]:
                return self._make_result('Bellman-Ford', 'O(V × E)', 'O(V)', error='Negative cycle detected')

        return self._make_result('Bellman-Ford', 'O(V × E)', 'O(V)',
            distances=distances, parent=parent, vertices=V, edges=graph.edge_count)

    # ==================== FLOYD-WARSHALL ====================
    @profile_algorithm
    def floyd_warshall(self, graph: Graph) -> Dict:
        self.reset()
        n = max(graph.vertices) + 1 if graph.vertices else 0
        dist = [[float('inf')] * n for _ in range(n)]
        next_v = [[None] * n for _ in range(n)]
        for i in range(n):
            dist[i][i] = 0
        for u in graph.vertices:
            for v, w in graph.get_neighbors(u):
                weight = w if w is not None else 1
                dist[u][v] = weight
                next_v[u][v] = v

        for k in range(n):
            for i in range(n):
                for j in range(n):
                    self.iterations += 1
                    self.comparisons += 1
                    if dist[i][k] + dist[k][j] < dist[i][j]:
                        dist[i][j] = dist[i][k] + dist[k][j]
                        next_v[i][j] = next_v[i][k]

        return self._make_result('Floyd-Warshall', 'O(V³)', 'O(V²)',
            distances=dist, next_vertex=next_v, vertices=len(graph.vertices), edges=graph.edge_count)

    # ==================== PRIM MST ====================
    @profile_algorithm
    def prim_mst(self, graph: Graph) -> Dict:
        self.reset()
        if graph.directed:
            return self._make_result("Prim's MST", 'O((V+E) log V)', 'O(V)', error='Graph must be undirected')
        start = min(graph.vertices)
        visited = {start}
        mst_edges = []
        total_weight = 0
        pq = []
        for neighbor, weight in graph.get_neighbors(start):
            w = weight if weight is not None else 1
            heapq.heappush(pq, (w, start, neighbor))

        while pq and len(visited) < len(graph.vertices):
            self.iterations += 1
            weight, u, v = heapq.heappop(pq)
            if v in visited:
                continue
            visited.add(v)
            mst_edges.append((u, v, weight))
            total_weight += weight
            for neighbor, w in graph.get_neighbors(v):
                edge_weight = w if w is not None else 1
                if neighbor not in visited:
                    heapq.heappush(pq, (edge_weight, v, neighbor))

        return self._make_result("Prim's MST", 'O((V+E) log V)', 'O(V)',
            mst_edges=mst_edges, total_weight=total_weight,
            vertices=len(graph.vertices), edges=graph.edge_count)

    # ==================== KRUSKAL MST ====================
    @profile_algorithm
    def kruskal_mst(self, graph: Graph) -> Dict:
        self.reset()
        if graph.directed:
            return self._make_result("Kruskal's MST", 'O(E log E)', 'O(V)', error='Graph must be undirected')
        parent = {v: v for v in graph.vertices}
        rank = {v: 0 for v in graph.vertices}

        def find(v):
            if parent[v] != v:
                parent[v] = find(parent[v])
            return parent[v]

        def union(u, v):
            ru, rv = find(u), find(v)
            if ru != rv:
                if rank[ru] < rank[rv]:
                    parent[ru] = rv
                elif rank[ru] > rank[rv]:
                    parent[rv] = ru
                else:
                    parent[rv] = ru
                    rank[ru] += 1

        edges = sorted(graph.edges, key=lambda x: x[2] if x[2] is not None else 1)
        mst_edges = []
        total_weight = 0
        for u, v, w in edges:
            weight = w if w is not None else 1
            if find(u) != find(v):
                union(u, v)
                mst_edges.append((u, v, weight))
                total_weight += weight
                if len(mst_edges) == len(graph.vertices) - 1:
                    break

        return self._make_result("Kruskal's MST", 'O(E log E)', 'O(V)',
            mst_edges=mst_edges, total_weight=total_weight,
            vertices=len(graph.vertices), edges=graph.edge_count)

    # ==================== LONGEST PATH DAG ====================
    @profile_algorithm
    def longest_path_dag(self, graph: Graph, start: int) -> Dict:
        self.reset()
        in_degree = {v: 0 for v in graph.vertices}
        for u in graph.vertices:
            for v, _ in graph.get_neighbors(u):
                in_degree[v] += 1
        queue = deque([v for v in graph.vertices if in_degree[v] == 0])
        topo_order = []
        while queue:
            self.iterations += 1
            vertex = queue.popleft()
            topo_order.append(vertex)
            for neighbor, _ in graph.get_neighbors(vertex):
                in_degree[neighbor] -= 1
                if in_degree[neighbor] == 0:
                    queue.append(neighbor)

        distances = {v: float('-inf') for v in graph.vertices}
        distances[start] = 0
        parent = {v: None for v in graph.vertices}
        for vertex in topo_order:
            if distances[vertex] != float('-inf'):
                for neighbor, weight in graph.get_neighbors(vertex):
                    w = weight if weight is not None else 1
                    if distances[vertex] + w > distances[neighbor]:
                        distances[neighbor] = distances[vertex] + w
                        parent[neighbor] = vertex

        return self._make_result('Longest Path (DAG)', 'O(V + E)', 'O(V)',
            distances=distances, parent=parent, topological_order=topo_order,
            vertices=len(graph.vertices), edges=graph.edge_count)

    # ==================== TOPOLOGICAL SORT (Kahn) ====================
    @profile_algorithm
    def topological_sort(self, graph: Graph) -> Dict:
        self.reset()
        in_degree = {v: 0 for v in graph.vertices}
        for u in graph.vertices:
            for v, _ in graph.get_neighbors(u):
                in_degree[v] += 1
        queue = deque([v for v in graph.vertices if in_degree[v] == 0])
        topo_order = []
        while queue:
            self.iterations += 1
            vertex = queue.popleft()
            topo_order.append(vertex)
            for neighbor, _ in graph.get_neighbors(vertex):
                in_degree[neighbor] -= 1
                if in_degree[neighbor] == 0:
                    queue.append(neighbor)

        valid = len(topo_order) == len(graph.vertices)
        return self._make_result('Topological Sort (Kahn)', 'O(V + E)', 'O(V)',
            topological_order=topo_order, valid=valid,
            vertices=len(graph.vertices), edges=graph.edge_count)

    # ==================== SCC (KOSARAJU) ====================
    @profile_algorithm
    def kosaraju_scc(self, graph: Graph) -> Dict:
        self.reset()
        if not graph.directed:
            return self._make_result('Kosaraju SCC', 'O(V + E)', 'O(V)', error='Graph must be directed')

        visited = set()
        finish_order = []

        def dfs1(v):
            self.recursions += 1
            visited.add(v)
            for neighbor, _ in graph.get_neighbors(v):
                if neighbor not in visited:
                    dfs1(neighbor)
            finish_order.append(v)

        for v in graph.vertices:
            if v not in visited:
                dfs1(v)

        # Build transpose
        transpose = Graph(directed=True)
        for v in graph.vertices:
            transpose.add_vertex(v)
        for u, v, w in graph.edges:
            transpose.add_edge(v, u, w if w else 1)

        visited.clear()
        sccs = []
        for v in reversed(finish_order):
            if v not in visited:
                stack = [v]
                component = []
                while stack:
                    node = stack.pop()
                    if node not in visited:
                        visited.add(node)
                        component.append(node)
                        for neighbor, _ in transpose.get_neighbors(node):
                            if neighbor not in visited:
                                stack.append(neighbor)
                sccs.append(component)

        return self._make_result('Kosaraju SCC', 'O(V + E)', 'O(V)',
            sccs=sccs, count=len(sccs), vertices=len(graph.vertices), edges=graph.edge_count)

    # ==================== A* SEARCH ====================
    @profile_algorithm
    def a_star(self, graph: Graph, start: int, target: int, heuristic: Optional[Dict[int, float]] = None) -> Dict:
        self.reset()
        if heuristic is None:
            heuristic = {v: 0 for v in graph.vertices}
        g_score = {v: float('inf') for v in graph.vertices}
        g_score[start] = 0
        f_score = {v: float('inf') for v in graph.vertices}
        f_score[start] = heuristic.get(start, 0)
        parent = {v: None for v in graph.vertices}
        open_set = [(f_score[start], start)]
        visited = set()

        while open_set:
            self.iterations += 1
            _, current = heapq.heappop(open_set)
            if current in visited:
                continue
            visited.add(current)
            if current == target:
                path = []
                curr = target
                while curr is not None:
                    path.append(curr)
                    curr = parent[curr]
                path.reverse()
                return self._make_result('A* Search', 'O((V+E) log V)', 'O(V)',
                    path=path, cost=g_score[target], found=True,
                    vertices=len(graph.vertices), edges=graph.edge_count)
            for neighbor, weight in graph.get_neighbors(current):
                self.accesses += 1
                self.comparisons += 1
                w = weight if weight is not None else 1
                tentative = g_score[current] + w
                if tentative < g_score[neighbor]:
                    parent[neighbor] = current
                    g_score[neighbor] = tentative
                    f_score[neighbor] = tentative + heuristic.get(neighbor, 0)
                    heapq.heappush(open_set, (f_score[neighbor], neighbor))

        return self._make_result('A* Search', 'O((V+E) log V)', 'O(V)',
            path=[], cost=float('inf'), found=False,
            vertices=len(graph.vertices), edges=graph.edge_count)