remove T_max and add a runtime for the code. where each path can be any length. and fix the collision flag
import gurobipy as gp
from gurobipy import GRB
import networkx as nx
import matplotlib.pyplot as plt

# Function to create a 2D grid graph with the specified number of rows and columns.
def create_grid_graph(n_rows, n_cols):
    G = nx.grid_2d_graph(n_rows, n_cols)  # Creates an undirected 2D grid graph
    G = G.to_directed()  # Converts it to a directed graph (for bidirectional movement)
    pos = {(i,j):(i,j) for i,j in G.nodes()}  # Assign positions for nodes for plotting purposes
    nx.set_node_attributes(G, pos, 'pos')  # Store the positions as node attributes
    return G

# Function to visualize the agents' paths on the graph using matplotlib.
def visualize_agents_paths(G, agents, paths):
    pos = nx.get_node_attributes(G, 'pos')  # Get node positions for plotting
    plt.figure(figsize=(8,8))
    nx.draw(G, pos, with_labels=False, node_size=50, node_color='lightgrey', arrows=False)  # Draw the grid

    colors = ['red', 'blue', 'green', 'orange', 'purple', 'brown']  # Colors for agents
    for idx, a in enumerate(agents):
        path = paths[a]
        # Draw the path traversed by each agent
        edges_in_path = list(zip(path, path[1:]))
        nx.draw_networkx_edges(G, pos, edgelist=edges_in_path, edge_color=colors[idx % len(colors)], width=2)
        # Highlight the starting position of each agent
        nx.draw_networkx_nodes(G, pos, nodelist=[agents[a]['start']], node_color=colors[idx % len(colors)], node_size=200, label=f"Agent {a} Start")
        # Highlight the goal positions of each agent
        nx.draw_networkx_nodes(G, pos, nodelist=agents[a]['goals'], node_color=colors[idx % len(colors)], node_shape='s', node_size=100)
    plt.legend()
    plt.show()

# Parameters
n_rows = 3  # Number of rows in the grid
n_cols = 3  # Number of columns in the grid
allow_collisions = True  # Whether to allow agents to collide

# Create the grid graph with the given dimensions
G = create_grid_graph(n_rows, n_cols)

# Define agents and their starting points and goals
agents = {
    1: {'start': (0, 0), 'goals': [(0, 0), (1, 2)]},  # Agent 1 starts at (0,0) and has goals to visit (0,0) and (1,2)
    2: {'start': (0, 1), 'goals': [(0, 2), (2, 0)]}   # Agent 2 starts at (0,1) and has goals to visit (0,2) and (2,0)
}

# Define the time horizon for the problem (how many time steps to consider)
T_max = 10  # Maximum number of time steps

# Create the Gurobi model
model = gp.Model("Multi-Agent Path Finding")

# Decision Variables
# x[a, i, t]: Binary variable, 1 if agent 'a' is at node 'i' at time 't'
x = {}
for a in agents:
    for t in range(T_max):
        for i in G.nodes():
            x[a, i, t] = model.addVar(vtype=GRB.BINARY, name=f"x_{a}{i}{t}")

# Movement variables
# m[a, i, j, t]: Binary variable, 1 if agent 'a' moves from node 'i' to node 'j' between time 't' and 't+1'
m = {}
for a in agents:
    for t in range(T_max - 1):
        for i in G.nodes():
            for j in G.successors(i):  # Successors are the nodes connected to i
                m[a, i, j, t] = model.addVar(vtype=GRB.BINARY, name=f"m_{a}{i}{j}{t}")
            # Include staying in place as a valid movement
            m[a, i, i, t] = model.addVar(vtype=GRB.BINARY, name=f"m{a}{i}{i}{t}")

# Constraints

# Ensure each agent starts at its designated start node at time 0
for a in agents:
    start_node = agents[a]['start']
    model.addConstr(x[a, start_node, 0] == 1, name=f"start{a}")

# Movement constraints for agents
for a in agents:
    for t in range(T_max - 1):
        for i in G.nodes():
            neighbors = list(G.successors(i)) + [i]  # Include staying in place
            # Agent can move only if it's at node i at time t
            model.addConstr(
                x[a, i, t] - gp.quicksum(x[a, j, t+1] for j in neighbors) <= 0,
                name=f"move_{a}{i}{t}"
            )
            # Link movement variables to position variables
            model.addConstr(
                m[a, i, i, t] + gp.quicksum(m[a, i, j, t] for j in G.successors(i)) == x[a, i, t],
                name=f"move_link_{a}{i}{t}"
            )
            # Ensure if an agent moves from i to j, they must be at i at time t and at j at time t+1
            for j in neighbors:
                model.addConstr(
                    m[a, i, j, t] <= x[a, i, t],
                    name=f"move_from_{a}{i}{j}{t}"
                )
                model.addConstr(
                    m[a, i, j, t] <= x[a, j, t+1],
                    name=f"move_to_{a}{i}{j}{t}"
                )

# Ensure each agent is at exactly one node at every time step
for a in agents:
    for t in range(T_max):
        model.addConstr(
            gp.quicksum(x[a, i, t] for i in G.nodes()) == 1,
            name=f"loc{a}{t}"
        )

# Collision avoidance constraints if collisions are not allowed
if not allow_collisions:
    # Ensure no two agents occupy the same node at the same time
    for t in range(T_max):
        for i in G.nodes():
            model.addConstr(
                gp.quicksum(x[a, i, t] for a in agents) <= 1,
                name=f"vertex_collision{i}{t}"
            )
    # Ensure no two agents swap edges at the same time
    for t in range(T_max - 1):
        for i in G.nodes():
            for j in G.successors(i):
                model.addConstr(
                    gp.quicksum(m[a, i, j, t] + m[a, j, i, t] for a in agents) <= 1,
                    name=f"edge_collision{i}{j}{t}"
                )

# Ensure each agent visits all its goal nodes at least once
for a in agents:
    for goal_node in agents[a]['goals']:
        model.addConstr(
            gp.quicksum(x[a, goal_node, t] for t in range(T_max)) >= 1,
            name=f"visit_{a}_{goal_node}"
        )

# Objective: Minimize total number of moves across all agents
model.setObjective(
    gp.quicksum(m[a, i, j, t] for a in agents for t in range(T_max - 1) 
                for i in G.nodes() for j in G.neighbors(i) if j != i),
    GRB.MINIMIZE
)

# Solve the model
model.optimize()

# Extract paths for each agent after solving
paths = {a: [] for a in agents}
for a in agents:
    for t in range(T_max):
        for i in G.nodes():
            if x[a,i,t].X > 0.5:  # If the variable is set to 1, the agent is at node i at time t
                paths[a].append(i)
                break  # Agent can only be at one node at a time

# Print the paths for each agent
for a in agents:
    print(f"Agent {a} path:")
    print(paths[a])

# Visualize the paths of the agents
visualize_agents_paths(G, agents, paths)