"""ch13_ant_colony.py
Companion code for Chapter 13: Optimization and Data Fitting.

Ant colony search for a short closed tour.
Edit the example settings in main(); numerical routines are reusable on import.
Run: python3 ch13_ant_colony.py [--no-plot]
Requires NumPy; optional plots require Matplotlib.
Fixed seeds reproduce runs with the same software; languages need not share sequences.
"""

import sys
import numpy as np


def main():
    make_plot = "--no-plot" not in sys.argv
    seed = 7
    rng = np.random.default_rng(seed)
    # Alternative: comment out the previous line and use fresh system entropy.
    # rng = np.random.default_rng()
    xy = np.array([[.10,.20],[.85,.15],[.65,.75],[.10,.70],[.55,.70],[.60,.35],
                   [.20,.55],[.80,.50],[.42,.12],[.50,.50],[.35,.35],[.40,.60]])
    D = np.linalg.norm(xy[:,None,:]-xy[None,:,:],axis=2)
    n_ants,max_iter,alpha,beta = 96,300,1.0,3.0
    rho,Q,tau0,elite_weight = 0.2,1.0,0.01,4.0
    tour,length,info = ant_colony(D,rng,n_ants,max_iter,alpha,beta,rho,Q,tau0,elite_weight)
    closed = np.r_[tour,tour[0]]
    print("Closed tour (city labels start at 1):", " -> ".join(str(i+1) for i in closed))
    print(f"Best length: {length:.12g}");print(info)
    print("Repeat with other seeds; global optimality is not guaranteed.")
    if make_plot:
        import matplotlib.pyplot as plt
        plt.plot(xy[closed,0],xy[closed,1],"k-o")
        for i,point in enumerate(xy): plt.annotate(str(i+1),point,xytext=(4,4),textcoords="offset points")
        plt.xlabel("x");plt.ylabel("y");plt.title("Ant colony closed tour",fontweight="normal")
        plt.axis("equal");plt.margins(.12);plt.grid(True);plt.tight_layout();plt.show()


def vector(x, name="data"):
    if np.iscomplexobj(x):
        raise ValueError(name + " must be real")
    a = np.asarray(x, dtype=float)
    if a.ndim != 1 or a.size == 0 or not np.all(np.isfinite(a)):
        raise ValueError(name + " must be a nonempty finite vector")
    return a


def checked(fun, x):
    value = np.asarray(fun(x))
    if value.ndim != 0 or np.iscomplexobj(value) or not np.isfinite(value):
        raise ValueError("Objective must return a finite real scalar")
    return float(value)


def bounds(lower, upper):
    lo, hi = vector(lower, "lower"), vector(upper, "upper")
    if lo.shape != hi.shape or not np.all(np.isfinite(hi - lo)) or np.any(hi <= lo):
        raise ValueError("Require matching finite bounds with upper > lower")
    return lo, hi


def integer(value, minimum, name):
    if not np.isscalar(value) or not np.isfinite(value) or int(value) != value or value < minimum:
        raise ValueError(name + " must be an integer >= " + str(minimum))
    return int(value)


def positive(value, name):
    if not np.isscalar(value) or not np.isfinite(value) or value <= 0:
        raise ValueError(name + " must be positive and finite")
    return float(value)

def ant_colony(D, rng, n_ants=96, max_iter=300, alpha=1.0, beta=3.0,
               rho=0.2, Q=1.0, tau0=0.01, elite_weight=4.0):
    """Symmetric TSP; log-scaled transition weights and full closed-tour deposits.

    Returns zero-based city indices, rotated to start at city zero.
    """
    if np.iscomplexobj(D):
        raise ValueError("Distances must be real")
    D = np.asarray(D, dtype=float); n = len(D)
    if D.ndim != 2 or D.shape != (n,n) or n < 3 or not np.all(np.isfinite(D)):
        raise ValueError("Require a finite square distance matrix with >=3 cities")
    off = ~np.eye(n, dtype=bool)
    if not np.allclose(D, D.T, rtol=1e-13, atol=0) or np.any(np.diag(D) != 0) or np.any(D[off] <= 0):
        raise ValueError("Require symmetric distances, zero diagonal, positive off-diagonal")
    M = integer(n_ants, 1, "n_ants"); K = integer(max_iter, 1, "max_iter")
    if not np.isfinite(alpha+beta+elite_weight) or min(alpha,beta,elite_weight) < 0 or not 0 < rho < 1:
        raise ValueError("Invalid pheromone settings")
    positive(Q,"Q"); positive(tau0,"tau0")
    tau = np.full((n,n), tau0); best = None; fb = np.inf
    def deposit(route, amount):
        for a,b in zip(route, np.roll(route,-1)):
            tau[a,b] += amount; tau[b,a] += amount
    for _ in range(K):
        tours, lengths = [], []
        for _ in range(M):
            route = [int(rng.integers(n))]; available = list(range(n)); available.remove(route[0])
            while available:
                j = route[-1]; ids = np.array(available)
                logw = alpha*np.log(tau[j,ids]) - beta*np.log(D[j,ids])
                w = np.exp(logw - np.max(logw)); w /= np.sum(w)
                city = int(rng.choice(ids, p=w)); route.append(city); available.remove(city)
            route = np.array(route); length = float(np.sum(D[route,np.roll(route,-1)]))
            tours.append(route); lengths.append(length)
            if length < fb:
                best, fb = route.copy(), length
        tau *= 1-rho
        for route,length in zip(tours,lengths):
            deposit(route, Q/length)
        deposit(best, elite_weight*Q/fb)
        tau = np.maximum(tau, np.finfo(float).tiny)
        if not np.all(np.isfinite(tau)):
            raise ValueError("Pheromone overflow; rescale distances or deposits")
    best = np.roll(best, -int(np.flatnonzero(best == 0)[0]))
    return best, fb, dict(iterations=K, evaluations=M*K, status="Search budget completed.")

def report(p, value, info):
    for j,x in enumerate(np.atleast_1d(p),1):
        print(f"p({j}) = {x:.12g}")
    print(f"Objective value: {value:.12g}")
    for key,value in info.items():
        print(f"{key}: {value}")


if __name__ == "__main__":
    main()