import gurobipy as gp
from gurobipy import GRB
import networkx as nx
import matplotlib.pyplot as plt

def create_grid_graph(n_rows, n_cols):
    G = nx.grid_2d_graph(n_rows, n_cols)
    G = G.to_directed()
    pos = {(i,j):(i,j) for i,j in G.nodes()}
    nx.set_node_attributes(G, pos, 'pos')
    return G

def compute_shortest_paths(G, required_nodes):
    import itertools
    distances = {}
    paths = {}
    for i, j in itertools.permutations(required_nodes, 2):
        length, path = nx.single_source_dijkstra(G, i, j)
        distances[i,j] = length
        paths[i,j] = path
    return distances, paths

def create_complete_graph(required_nodes, distances):
    G_complete = nx.DiGraph()
    for i in required_nodes:
        G_complete.add_node(i)
    for (i,j), length in distances.items():
        G_complete.add_edge(i, j, weight=length)
    return G_complete

def solve_tsp_complete_graph(G_complete, start_node):
    model = gp.Model("TSP")
    nodes = list(G_complete.nodes)
    edges = list(G_complete.edges)
    num_nodes = len(nodes)
    
    x = {}
    for i, j in edges:
        x[i, j] = model.addVar(vtype=GRB.BINARY, name=f"x_{i}{j}")
    
    model.setObjective(gp.quicksum(G_complete[i][j]['weight'] * x[i,j] for i,j in edges), GRB.MINIMIZE)
    
    for i in nodes:
        model.addConstr(gp.quicksum(x[i,j] for j in nodes if (i,j) in x) == 1, name=f"out{i}")
        model.addConstr(gp.quicksum(x[j,i] for j in nodes if (j,i) in x) == 1, name=f"in_{i}")
    
    u = {}
    for i in nodes:
        u[i] = model.addVar(vtype=GRB.CONTINUOUS, lb=0, ub=num_nodes, name=f"u_{i}")
    
    model.addConstr(u[start_node] == 0, name='subtour_elimination_start')
    
    for i, j in edges:
        if i != start_node and j != start_node:
            model.addConstr(u[i] - u[j] + num_nodes * x[i,j] <= num_nodes - 1, name=f"subtour_elimination_{i}_{j}")
    
    model.optimize()
    
    solution = []
    for i, j in edges:
        if x[i,j].X > 0.5:
            solution.append((i,j))
    return solution

def reconstruct_full_path(tour, paths):
    full_path = []
    for i, j in tour:
        path = paths[i,j][:-1]
        full_path.extend(path)
    full_path.append(tour[-1][1])
    return full_path

def visualize_solution(G, full_paths, required_nodes_list):
    pos = nx.get_node_attributes(G, 'pos')
    plt.figure(figsize=(10,10))
    nx.draw(G, pos, with_labels=False, node_size=50, node_color='lightgrey', arrows=False)
    
    colors = ['red', 'blue', 'green', 'yellow', 'purple']  # Add more colors if needed
    for i, (full_path, required_nodes) in enumerate(zip(full_paths, required_nodes_list)):
        color = colors[i % len(colors)]
        nx.draw_networkx_nodes(G, pos, nodelist=required_nodes, node_color=color, node_size=100)
        edges_in_path = list(zip(full_path, full_path[1:]))
        nx.draw_networkx_edges(G, pos, edgelist=edges_in_path, edge_color=color, width=2, arrows=True)
    
    plt.show()

def solve_multi_agent_tsp(G, agents_required_nodes):
    full_paths = []
    for required_nodes in agents_required_nodes:
        distances, paths = compute_shortest_paths(G, required_nodes)
        G_complete = create_complete_graph(required_nodes, distances)
        start_node = required_nodes[0]  # Assume first node is the start/home node
        solution_edges = solve_tsp_complete_graph(G_complete, start_node)
        
        tour = []
        current_node = start_node
        visited = set()
        while len(visited) < len(required_nodes):
            for i,j in solution_edges:
                if i == current_node and j not in visited:
                    tour.append((i,j))
                    visited.add(i)
                    current_node = j
                    break
            else:
                break
        
        # Add return to start node
        tour.append((current_node, start_node))
        
        full_path = reconstruct_full_path(tour, paths)
        full_paths.append(full_path)
    
    return full_paths

# Example usage
n_rows, n_cols = 3, 3
G = create_grid_graph(n_rows, n_cols)

# Define required nodes for two agents
agents_required_nodes = [
    [(0,0), (1, 0), (1, 2)],  # Agent 1
    [(0,1), (0, 2), (2, 0)]   # Agent 2
]


# Solve TSP for multiple agents
full_paths = solve_multi_agent_tsp(G, agents_required_nodes)

# Print results
for i, path in enumerate(full_paths):
    print(f"Optimal path for Agent {i+1}: {path}")

# Visualize the solution
visualize_solution(G, full_paths, agents_required_nodes)

i want to check for collisions in the result solution. and i want better visualization where the lines of the agents path wont be overlaying but with a small offset to see them. and to identify the target nodes of each agent and the start point too. also add nodes index on the grid like (0,0) and so on