% ch13_ant_colony.m
% Companion code for Chapter 13: Optimization and Data Fitting.
%
% Ant colony optimization for the symmetric traveling-salesman problem.
% Find a short closed route visiting every city once and returning to the
% starting city. Change the city coordinates and search settings below.
% Distances are Euclidean, in the same units as the coordinates.
% This is a heuristic: the best route found need not be the shortest possible.
% Repeat with different seeds and search budgets to assess reliability.

clear; close all; clc;

%% City coordinates (one city per row)
XY = [0.10 0.20
      0.85 0.15
      0.65 0.75
      0.10 0.70
      0.55 0.70
      0.60 0.35
      0.20 0.55
      0.80 0.50
      0.42 0.12
      0.50 0.50
      0.35 0.35
      0.40 0.60];

%% Search settings
n_ants = 96;
max_iter = 300;
alpha = 1.0;              % Importance of pheromone
beta = 3.0;               % Importance of short edges (inverse distance)
rho = 0.20;               % Fraction of pheromone evaporated each iteration
Q = 1.0;                  % Deposit scale: each ant deposits Q / route length
tau0 = 0.01;              % Initial pheromone on each edge
elite_weight = 4.0;       % Extra deposit for best route; 0 disables it
seed = 7;                 % Fixed seed for reproducible runs
make_plot = true;

rng(seed, 'twister');
% Alternative: comment out the line above and uncomment the line below
% to initialize from the current time.
% rng('shuffle');

%% Distances and search
validateattributes(XY, {'double'}, {'real','2d','finite','ncols',2});
n = size(XY, 1);
if n < 3
    error('Supply at least three distinct cities.');
end
D = zeros(n);
for i = 1:n
    for j = i+1:n
        D(i,j) = hypot(XY(i,1)-XY(j,1), XY(i,2)-XY(j,2));
        D(j,i) = D(i,j);
    end
end
[tour, best_length, info] = ant_colony(D, n_ants, max_iter, ...
    alpha, beta, rho, Q, tau0, elite_weight);

%% Results
closed_tour = [tour tour(1)];
fprintf('\nAnt colony optimization\n');
fprintf('Cities:             %d\n', n);
fprintf('Best route length:  %.12g\n', best_length);
fprintf('Iterations:         %d\n', info.iterations);
fprintf('Routes evaluated:   %d\n', info.evaluations);
fprintf('Best closed route (starting at city 1):\n');
fprintf('%d ', closed_tour);
fprintf('\nStatus: Search budget completed.\n');
fprintf('Repeat with other seeds; global optimality is not guaranteed.\n');

%% Optional route plot
if make_plot
    figure('Color', 'w');
    plot(XY(closed_tour,1), XY(closed_tour,2), 'k-o', ...
        'LineWidth', 1.5, 'MarkerFaceColor', 'w', 'MarkerSize', 7);
    hold on;
    plot(XY(1,1), XY(1,2), 'ks', 'MarkerFaceColor', 'k', 'MarkerSize', 7);
    span = max(XY,[],1) - min(XY,[],1);
    pad = 0.08*max(span);
    for i = 1:n
        text(XY(i,1)+pad/4, XY(i,2)+pad/4, sprintf('%d',i), 'FontSize',11);
    end
    xlabel('x', 'FontSize',14); ylabel('y', 'FontSize',14);
    title('Ant colony optimization: best route found', ...
        'FontSize',14,'FontWeight','normal');
    set(gca,'FontSize',14);
    axis equal;
    xlim([min(XY(:,1))-pad, max(XY(:,1))+pad]);
    ylim([min(XY(:,2))-pad, max(XY(:,2))+pad]);
    grid on; box on;
end

%% Reusable numerical routine
function [best_tour, best_length, info] = ant_colony(D, M, K, ...
        alpha, beta, rho, Q, tau0, elite_weight)
% D is a symmetric distance matrix with zero diagonal and positive off-diagonal
% entries. The routine also accepts non-Euclidean symmetric distance matrices.
% Pheromone is updated only after all ants have completed their routes.
% Copy this routine and its helper to ant_colony.m for reuse.

    validateattributes(D, {'double'}, {'real','2d','square','finite','nonempty'});
    n = size(D,1);
    offdiag = ~eye(n);
    if n < 3 || any(diag(D) ~= 0) || any(D(offdiag) <= 0) || ~isequal(D,D.')
        error('Require symmetric D, zero diagonal, and positive distances between distinct cities.');
    end
    validateattributes(M, {'double'}, {'scalar','finite','integer','positive'});
    validateattributes(K, {'double'}, {'scalar','finite','integer','positive'});
    validateattributes(alpha, {'double'}, {'scalar','finite','nonnegative'});
    validateattributes(beta, {'double'}, {'scalar','finite','nonnegative'});
    validateattributes(rho, {'double'}, {'scalar','finite','>',0,'<',1});
    validateattributes(Q, {'double'}, {'scalar','finite','positive'});
    validateattributes(tau0, {'double'}, {'scalar','finite','positive'});
    validateattributes(elite_weight, {'double'}, {'scalar','finite','nonnegative'});

    tau = tau0*ones(n);
    best_length = inf;
    best_tour = [];
    for k = 1:K
        tours = zeros(M,n);
        lengths = zeros(M,1);
        for m = 1:M
            route = zeros(1,n);
            route(1) = randi(n);
            visited = false(1,n);
            visited(route(1)) = true;
            for step = 2:n
                current = route(step-1);
                available = find(~visited);
                % Relative weights proportional to tau^alpha / distance^beta.
                % Logarithms avoid underflow in very small relative weights.
                logw = alpha*log(tau(current,available)) ...
                       - beta*log(D(current,available));
                weights = exp(logw - max(logw));
                cumulative = cumsum(weights/sum(weights));
                cumulative(end) = 1;  % Protect against rounding at the endpoint.
                j = find(rand < cumulative, 1, 'first');
                route(step) = available(j);
                visited(route(step)) = true;
            end
            tours(m,:) = route;
            lengths(m) = sum(D(sub2ind([n n], route, route([2:end 1]))));
        end
        [length_k, idx] = min(lengths);
        if length_k < best_length
            best_length = length_k;
            best_tour = tours(idx,:);
        end

        tau = (1-rho)*tau;
        for m = 1:M
            tau = deposit(tau, tours(m,:), Q/lengths(m));
        end
        tau = deposit(tau, best_tour, elite_weight*Q/best_length);
        tau = max(tau, realmin);  % Keep logarithms finite after evaporation.
        if any(~isfinite(tau(:)))
            error('Pheromone overflow: rescale distances or deposit settings.');
        end
    end
    % Rotate the route for display without changing its length.
    j = find(best_tour == 1, 1);
    best_tour = best_tour([j:end 1:j-1]);
    info = struct('iterations', K, 'evaluations', M*K);
end

function tau = deposit(tau, route, amount)
% Include the closing edge and deposit on both directions of each edge.
    n = numel(route);
    for s = 1:n
        i = route(s);
        j = route(mod(s,n)+1);
        tau(i,j) = tau(i,j) + amount;
        tau(j,i) = tau(j,i) + amount;
    end
end