ants.py

import os
import numpy as np
import networkx as nx
import matplotlib
matplotlib.use('Agg') 
import matplotlib.pyplot as plt
from scipy.sparse.csgraph import shortest_path

# Import visualizer's sorter to ensure coordinates match the app perfectly
from visualizer import get_sorted_nodes

def get_physical_path(start, end, preds):
    path = []
    curr = end
    while curr != start and curr >= 0:
        path.append(curr)
        curr = preds[start, curr]
    path.reverse()
    return path

def run_ant_colony(distances, coords, n_ants=50, n_iterations=100, decay=0.1, alpha=1.0, beta=2.0):
    n_nodes = distances.shape[0]
    
    dist_matrix_for_pathing = np.where(distances == 0, np.inf, distances)
    np.fill_diagonal(dist_matrix_for_pathing, 0)
    all_pairs_distances, predecessors = shortest_path(csgraph=dist_matrix_for_pathing, directed=False, return_predecessors=True)
    
    pheromones = np.ones((n_nodes, n_nodes)) * 0.1
    best_macro_tour = None
    best_length = float('inf')
    
    start_node = np.lexsort((coords[:, 0], coords[:, 1]))[0]
    
    for iteration in range(n_iterations):
        all_tours = []
        all_lengths = []
        
        for ant in range(n_ants):
            unvisited = set(range(n_nodes))
            current_node = start_node 
            tour = [current_node]
            unvisited.remove(current_node)
            tour_length = 0.0
            
            while unvisited:
                candidates = list(unvisited)
                pher_values = pheromones[current_node, candidates]
                dist_values = all_pairs_distances[current_node, candidates]
                heuristic = 1.0 / (dist_values + 1e-10) 
                
                probabilities = (pher_values ** alpha) * (heuristic ** beta)
                if probabilities.sum() == 0:
                    probabilities = np.ones(len(candidates)) / len(candidates)
                else:
                    probabilities /= probabilities.sum() 
                
                next_node = np.random.choice(candidates, p=probabilities)
                
                tour.append(next_node)
                tour_length += all_pairs_distances[current_node, next_node]
                unvisited.remove(next_node)
                current_node = next_node
                
            tour_length += all_pairs_distances[tour[-1], tour[0]]
            tour.append(tour[0])
            
            all_tours.append(tour)
            all_lengths.append(tour_length)
            
            if tour_length < best_length:
                best_length = tour_length
                best_macro_tour = tour
                
        pheromones *= (1.0 - decay) 
        
        for tour, length in zip(all_tours, all_lengths):
            deposit_amount = 100.0 / length 
            for i in range(len(tour) - 1):
                u, v = tour[i], tour[i+1]
                pheromones[u, v] += deposit_amount
                pheromones[v, u] += deposit_amount 
            
    # Unpack the macro tour into physical steps
    physical_tour = [best_macro_tour[0]]
    for i in range(len(best_macro_tour) - 1):
        start_n = best_macro_tour[i]
        end_n = best_macro_tour[i+1]
        segment = get_physical_path(start_n, end_n, predecessors)
        physical_tour.extend(segment)
        
    return physical_tour, best_length

def draw_base_graph_edges(ax, distances, coords, color='red'):
    n_nodes = distances.shape[0]
    for u in range(n_nodes):
        for v in range(u + 1, n_nodes): 
            if distances[u, v] > 0 and distances[u, v] != np.inf:
                start_c = coords[u]
                end_c = coords[v]
                ax.annotate("", xy=end_c, xytext=start_c, 
                            arrowprops=dict(arrowstyle="-", color=color, linewidth=4.0, alpha=0.6, zorder=1))

def visualize_tour_app(coords, physical_tour, title, distances, width, height):
    fig, ax = plt.figure(figsize=(10, 10)), plt.gca()
    
    # Force grid to match app dimensions
    ax.set_xlim(-0.5, width - 0.5)
    ax.set_ylim(-0.5, height - 0.5)
    
    xs = coords[:, 0]
    ys = coords[:, 1]
    
    draw_base_graph_edges(ax, distances, coords)
    ax.scatter(xs, ys, c='blue', s=100, zorder=5)
    
    entrance_idx = physical_tour[0]
    ax.scatter(coords[entrance_idx, 0], coords[entrance_idx, 1], c='yellow', edgecolors='black', s=400, marker='*', zorder=10, label="Entrance")
    
    tour_coords = coords[physical_tour]
    
    for i in range(len(tour_coords) - 1):
        start_c = tour_coords[i]
        end_c = tour_coords[i+1]
        
        dx = end_c[0] - start_c[0]
        dy = end_c[1] - start_c[1]
        
        ax.plot([start_c[0], end_c[0]], [start_c[1], end_c[1]], 
                color="blue", linewidth=2.0, linestyle="--", zorder=6)
        
        mid_x = start_c[0] + dx * 0.50
        mid_y = start_c[1] + dy * 0.50
        target_x = start_c[0] + dx * 0.51
        target_y = start_c[1] + dy * 0.51
        
        ax.annotate("", xy=(target_x, target_y), xytext=(mid_x, mid_y),
                    arrowprops=dict(arrowstyle="-|>,head_width=0.4,head_length=0.8", 
                                    color="blue", linewidth=2.0, zorder=7))

    ax.set_title(title, pad=20, fontsize=14, fontweight='bold')
    ax.invert_yaxis() 
    ax.grid(True, linestyle=':', alpha=0.6)
    
    from matplotlib.lines import Line2D
    custom_lines = [Line2D([0], [0], color="blue", linewidth=2.0, linestyle="--"),
                    Line2D([0], [0], marker='*', color='w', markerfacecolor='yellow', markeredgecolor='black', markersize=15),
                    Line2D([0], [0], color="red", linewidth=4.0, alpha=0.6)]
    ax.legend(custom_lines, ['Worker Route', 'Entrance', 'Available Hallways'], loc="best")
    
    save_dir = "temp_visuals"
    os.makedirs(save_dir, exist_ok=True)
    save_filename = os.path.join(save_dir, "app_ant_route.png")
    
    plt.savefig(save_filename, bbox_inches='tight')
    plt.close()
    return save_filename

def solve_and_visualize(G, width, height):
    """Main function to be called from app.py"""
    sorted_nodes = get_sorted_nodes(G)
    n_nodes = len(sorted_nodes)
    
    # 1. Build Coordinates Matrix
    coords = np.array(sorted_nodes, dtype=np.float64)
    
    # 2. Build Distance Matrix (1 Move = 1 Travel Meter)
    distances = np.zeros((n_nodes, n_nodes), dtype=np.float64)
    for u, v in G.edges():
        idx_u = sorted_nodes.index(u)
        idx_v = sorted_nodes.index(v)
        distances[idx_u, idx_v] = 1.0
        distances[idx_v, idx_u] = 1.0
        
    # 3. Run the swarm
    physical_tour, best_length = run_ant_colony(distances, coords)
    
    # 4. Draw the image and return the path
    img_path = visualize_tour_app(coords, physical_tour, f"Ant Colony Optimized Route\nTotal Moves: {int(best_length)}", distances, width, height)
    return img_path, int(best_length)