algorithms.py

import heapq
import math
from typing import Dict, List, Tuple
from models import TransportationGraph, Edge

class DisjointSet:
    """Disjoint Set Union for Kruskal's algorithm."""
    def __init__(self, elements):
        self.parent = {e: e for e in elements}
        self.rank = {e: 0 for e in elements}
    
    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]
    
    def union(self, x, y):
        px, py = self.find(x), self.find(y)
        if px == py:
            return False
        if self.rank[px] < self.rank[py]:
            self.parent[px] = py
        elif self.rank[px] > self.rank[py]:
            self.parent[py] = px
        else:
            self.parent[py] = px
            self.rank[px] += 1
        return True

def kruskal_mst(graph: TransportationGraph, prioritize_critical: bool = True, include_new_roads: bool = True) -> Tuple[List[Dict], float]:
    """
    Kruskal's Minimum Spanning Tree algorithm.
    Returns: (mst_edges, total_cost)
    """
    # Filter edges based on include_new_roads flag
    if include_new_roads:
        all_edges = list(graph.edges.values())
    else:
        all_edges = [e for e in graph.edges.values() if not e.is_new]
    
    critical_types = {'Medical', 'Government', 'Airport', 'Transit Hub'}
    critical_nodes = {n.id for n in graph.nodes.values() if n.type in critical_types}
    
    def edge_cost(edge: Edge) -> float:
        if edge.is_new:
            base = edge.construction_cost
        else:
            base = edge.distance * 10
        
        if prioritize_critical:
            if edge.from_id in critical_nodes or edge.to_id in critical_nodes:
                base *= 0.5
        
        pop_from = graph.nodes[edge.from_id].population if edge.from_id in graph.nodes else 0
        pop_to = graph.nodes[edge.to_id].population if edge.to_id in graph.nodes else 0
        avg_pop = (pop_from + pop_to) / 2
        if avg_pop > 300000:
            base *= 0.7
        
        return base
    
    sorted_edges = sorted(all_edges, key=edge_cost)
    
    all_nodes = list(graph.nodes.keys())
    dsu = DisjointSet(all_nodes)
    
    mst_edges = []
    total_cost = 0.0
    
    for edge in sorted_edges:
        if dsu.union(edge.from_id, edge.to_id):
            mst_edges.append(edge)
            if edge.is_new:
                total_cost += edge.construction_cost
            else:
                total_cost += edge.distance * 10
            
            if len(mst_edges) == len(all_nodes) - 1:
                break
    
    # Convert to serializable format
    serializable_edges = [{
        'from_id': e.from_id,
        'to_id': e.to_id,
        'distance': e.distance,
        'is_new': e.is_new,
        'construction_cost': e.construction_cost,
        'condition': e.condition,
        'capacity': e.capacity
    } for e in mst_edges]
    
    return serializable_edges, total_cost

def dijkstra_shortest_path(graph: TransportationGraph, start: str, end: str, 
                          time_period: int = 0) -> Tuple[List[str], float, float]:
    """
    Dijkstra's algorithm for shortest path with time-dependent traffic.
    Returns: (path, travel_time_minutes, total_distance_km)
    """
    if start not in graph.nodes or end not in graph.nodes:
        return [], float('inf'), 0
    
    pq = [(0, start, [start], 0)]
    best_cost = {start: 0}
    
    while pq:
        current_cost, current_node, path, current_dist = heapq.heappop(pq)
        
        if current_cost > best_cost.get(current_node, float('inf')):
            continue
        
        if current_node == end:
            return path, current_cost, current_dist
        
        for neighbor, edge in graph.get_neighbors(current_node):
            travel_time = edge.get_travel_time(time_period)
            new_cost = current_cost + travel_time
            new_dist = current_dist + edge.distance
            
            if new_cost < best_cost.get(neighbor, float('inf')):
                best_cost[neighbor] = new_cost
                new_path = path + [neighbor]
                heapq.heappush(pq, (new_cost, neighbor, new_path, new_dist))
    
    return [], float('inf'), 0

def emergency_heuristic(graph: TransportationGraph, node_id: str, goal_id: str) -> float:
    """Heuristic for A* emergency routing."""
    node = graph.nodes.get(node_id)
    goal = graph.nodes.get(goal_id)
    if not node or not goal:
        return 0
    
    dx = abs(node.x - goal.x) 
    dy = abs(node.y - goal.y)
    
    distance_km = math.sqrt((dx * 111.32)**2 + (dy * 110.57)**2)
    
    return distance_km * 0.75

def a_star_emergency_routing(graph: TransportationGraph, start: str, end: str,
                            time_period: int = 0) -> Tuple[List[str], float, float]:
    """A* search algorithm for emergency vehicle routing."""
    if start not in graph.nodes:
        return [], float('inf'), 0
    if end not in graph.nodes:
        return [], float('inf'), 0
    
    if start == end:
        return [start], 0, 0
    
    counter = 0
    start_h = emergency_heuristic(graph, start, end)
    pq = [(start_h, counter, 0, start, [start], 0)]
    
    best_g = {start: 0}
    visited = set()
    
    while pq:
        f_score, _, current_g, current_node, path, current_dist = heapq.heappop(pq)
        
        if current_g > best_g.get(current_node, float('inf')):
            continue
        
        if current_node == end:
            return path, current_g, current_dist
        
        visited.add(current_node)
        
        for neighbor, edge in graph.get_neighbors(current_node):
            if neighbor in visited:
                continue
            
            base_time = edge.get_travel_time(time_period)
            condition_bonus = 0.6 + (edge.condition / 10) * 0.4
            travel_time = base_time * condition_bonus
            
            new_g = current_g + travel_time
            new_dist = current_dist + edge.distance
            
            if new_g < best_g.get(neighbor, float('inf')):
                best_g[neighbor] = new_g
                h = emergency_heuristic(graph, neighbor, end)
                f = new_g + h
                new_path = path + [neighbor]
                counter += 1
                heapq.heappush(pq, (f, counter, new_g, neighbor, new_path, new_dist))
    
    return [], float('inf'), 0

def dynamic_programming_bus_scheduling(graph: TransportationGraph, route_id: str, 
                                       available_buses: int) -> Dict:
    """Dynamic Programming for optimal bus scheduling."""
    if route_id not in graph.bus_routes:
        return {}
    
    route = graph.bus_routes[route_id]
    stops = route['stops']
    n = len(stops)
    
    populations = []
    for stop_id in stops:
        node = graph.nodes.get(stop_id)
        populations.append(node.population if node else 0)
    
    dp = [[0] * (available_buses + 1) for _ in range(n + 1)]
    
    for i in range(1, n + 1):
        for j in range(1, available_buses + 1):
            dp[i][j] = dp[i-1][j]
            
            for prev in range(i):
                segment_length = i - prev
                buses_needed = max(1, segment_length // 3)
                
                if buses_needed <= j:
                    coverage = sum(populations[prev:i])
                    if dp[prev][j - buses_needed] + coverage > dp[i][j]:
                        dp[i][j] = dp[prev][j - buses_needed] + coverage
    
    stop_details = []
    for stop_id in stops:
        node = graph.nodes.get(stop_id)
        stop_details.append({
            'id': stop_id,
            'name': node.name if node else stop_id,
            'population': populations[stops.index(stop_id)]
        })
    
    return {
        'route_id': route_id,
        'stops': stops,
        'stop_details': stop_details,
        'max_coverage': dp[n][available_buses],
        'total_population': sum(populations),
        'available_buses': available_buses,
        'coverage_percentage': round(dp[n][available_buses] / sum(populations) * 100, 1) if sum(populations) > 0 else 0
    }

def dp_road_maintenance_allocation(graph: TransportationGraph, budget: float) -> Dict:
    """DP for road maintenance allocation (0/1 knapsack)."""
    existing_roads = [e for e in graph.edges.values() if not e.is_new]
    
    road_data = []
    for edge in existing_roads:
        cost_to_repair = int((10 - edge.condition) * 15)
        if cost_to_repair <= 0:
            continue
            
        pop_from = graph.nodes[edge.from_id].population if edge.from_id in graph.nodes else 0
        pop_to = graph.nodes[edge.to_id].population if edge.to_id in graph.nodes else 0
        
        traffic_importance = edge.capacity / 1000
        population_importance = (pop_from + pop_to) / 100000
        benefit = int(traffic_importance * population_importance * (edge.condition + 1))
        
        road_data.append({
            'edge': edge,
            'cost': cost_to_repair,
            'benefit': max(1, benefit),
            'condition': edge.condition,
            'road': f"{edge.from_id}-{edge.to_id}",
            'from_name': graph.nodes[edge.from_id].name if edge.from_id in graph.nodes else edge.from_id,
            'to_name': graph.nodes[edge.to_id].name if edge.to_id in graph.nodes else edge.to_id,
            'capacity': edge.capacity
        })
    
    budget_int = int(budget)
    dp = [0] * (budget_int + 1)
    selected = [[] for _ in range(budget_int + 1)]
    
    for road in road_data:
        cost = road['cost']
        benefit = road['benefit']
        
        for b in range(budget_int, cost - 1, -1):
            if dp[b - cost] + benefit > dp[b]:
                dp[b] = dp[b - cost] + benefit
                selected[b] = selected[b - cost] + [road]
    
    best_b = budget_int
    while best_b > 0 and not selected[best_b]:
        best_b -= 1
    
    roads_to_repair = []
    for road in selected[best_b] if best_b > 0 else []:
        roads_to_repair.append({
            'road': road['road'],
            'from_name': road['from_name'],
            'to_name': road['to_name'],
            'cost': road['cost'],
            'benefit': road['benefit'],
            'condition': road['condition'],
            'capacity': road['capacity']
        })
    
    return {
        'budget': budget,
        'total_cost': sum(r['cost'] for r in selected[best_b]) if best_b > 0 else 0,
        'total_benefit': dp[best_b] if best_b > 0 else 0,
        'roads_to_repair': roads_to_repair,
        'num_roads_repaired': len(selected[best_b]) if best_b > 0 else 0
    }

def greedy_traffic_signal_optimization(graph: TransportationGraph, 
                                      intersection_id: str) -> Dict:
    """Greedy traffic signal optimization."""
    if intersection_id not in graph.nodes:
        return {}
    
    connected = []
    for neighbor, edge in graph.get_neighbors(intersection_id):
        pop = graph.nodes[neighbor].population if neighbor in graph.nodes else 0
        priority_score = edge.capacity * (pop + 1)
        connected.append({
            'neighbor': neighbor,
            'name': graph.nodes[neighbor].name if neighbor in graph.nodes else neighbor,
            'flow': edge.capacity,
            'population': pop,
            'priority_score': priority_score,
            'edge': edge
        })
    
    connected.sort(key=lambda x: x['priority_score'], reverse=True)
    
    total_flow = sum(c['flow'] for c in connected) or 1
    
    signal_plan = []
    for i, conn in enumerate(connected):
        proportion = conn['flow'] / total_flow
        green_time = max(15, min(45, int(60 * proportion * 0.8)))
        
        if i < len(connected) // 3:
            priority = 'HIGH'
            priority_icon = '🔴'
        elif i < 2 * len(connected) // 3:
            priority = 'MEDIUM'
            priority_icon = '🟡'
        else:
            priority = 'LOW'
            priority_icon = '🟢'
        
        signal_plan.append({
            'direction': f"{intersection_id} → {conn['neighbor']}",
            'direction_name': f"{graph.nodes[intersection_id].name} → {conn['name']}",
            'green_time_seconds': green_time,
            'priority': priority,
            'priority_icon': priority_icon,
            'flow': conn['flow']
        })
    
    return {
        'intersection': intersection_id,
        'intersection_name': graph.nodes[intersection_id].name,
        'signal_plan': signal_plan,
        'total_cycle_time': sum(s['green_time_seconds'] for s in signal_plan)
    }