import random
import matplotlib.pyplot as plt

cities = ['A','B','C','D']
dist = [[0,10,10,30],
        [10,0,40,10],
        [10,40,0,20],
        [30,10,20,0]]

n = len(cities)
ants, iterations = 10, 50
alpha, beta, rho, Q = 1, 2, 0.5, 100

pheromone = [[1]*n for _ in range(n)]
best, best_cost = None, float('inf')
history = []

def cost(path):
    return sum(dist[path[i]][path[(i+1)%n]] for i in range(n))

for _ in range(iterations):
    paths = []

    for _ in range(ants):
        path = [random.randrange(n)]

        while len(path) < n:
            i = path[-1]
            choices = [j for j in range(n) if j not in path]
            weights = [pheromone[i][j]**alpha *
                       (1/dist[i][j])**beta for j in choices]
            path.append(random.choices(choices, weights=weights)[0])

        paths.append(path)
        d = cost(path)

        if d < best_cost:
            best, best_cost = path.copy(), d

    for i in range(n):
        for j in range(n):
            pheromone[i][j] *= (1-rho)

    for path in paths:
        for i in range(n):
            a, b = path[i], path[(i+1)%n]
            pheromone[a][b] += Q/cost(path)
            pheromone[b][a] += Q/cost(path)

    history.append(best_cost)

print("Best Path:", " -> ".join(cities[i] for i in best+[best[0]]))
print("Best Distance:", best_cost)

plt.plot(history)
plt.xlabel("Iterations")
plt.ylabel("Best Distance")
plt.title("ACO Convergence")
plt.grid()
plt.show()
