// ch13_ant_colony.cpp
// Companion code for Chapter 13: Optimization and Data Fitting.
// Standalone C++11; standard library only; console output without plotting.
// Compile: c++ -std=c++11 -O2 -Wall -Wextra -pedantic ch13_ant_colony.cpp -o ch13_ant_colony
// Edit the example settings in main(); routines above main() can be reused.
// Fixed seeds reproduce runs with the same standard-library implementation.
// Random distributions need not give identical sequences across compilers/languages.

#include <algorithm>
#include <array>
#include <cmath>
#include <functional>
#include <iomanip>
#include <iostream>
#include <limits>
#include <numeric>
#include <random>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

using Vec = std::vector<double>;
using Mat = std::vector<Vec>;
using Function = std::function<double(const Vec&)>;
using ScalarFunction = std::function<double(double)>;
const double pi = std::acos(-1.0);

void require(bool condition, const std::string& message) {
    if (!condition) throw std::invalid_argument(message);
}
void finite_vector(const Vec& x) {
    require(!x.empty(), "A nonempty vector is required.");
    for (double v : x) require(std::isfinite(v), "Data must be finite.");
}
void check_bounds(const Vec& lo, const Vec& hi) {
    finite_vector(lo); finite_vector(hi);
    require(lo.size() == hi.size(), "Bounds must have matching dimensions.");
    for (std::size_t j=0; j<lo.size(); ++j)
        require(hi[j]>lo[j] && std::isfinite(hi[j]-lo[j]), "Require finite upper > lower.");
}
double objective_value(double value) {
    require(std::isfinite(value), "Objective must return a finite real scalar.");
    return value;
}
double clip(double x, double lo, double hi) { return std::max(lo,std::min(hi,x)); }
Vec linspace(double a, double b, int n) {
    require(n>=2, "Grid needs at least two points.");
    Vec x(n);
    for (int i=0; i<n; ++i) x[i]=a+(b-a)*i/(n-1);
    return x;
}
struct Random {
    std::mt19937 engine;
    std::uniform_real_distribution<double> uniform{0.0,1.0};
    std::normal_distribution<double> normal{0.0,1.0};
    explicit Random(unsigned int seed) : engine(seed) {}
    double u() { return uniform(engine); }
    double z() { return normal(engine); }
    int index(int n) { return std::uniform_int_distribution<int>(0,n-1)(engine); }
};
struct Result {
    Vec point;
    double value=0.0;
    long long evaluations=0;
    int iterations=0;
    std::string status;
};
void print_result(const Result& r) {
    std::cout << std::setprecision(12);
    for (std::size_t j=0; j<r.point.size(); ++j)
        std::cout << "p(" << j+1 << ") = " << r.point[j] << '\n';
    std::cout << "Objective value: " << r.value << "\nIterations/updates: " << r.iterations
              << "\nFunction calls: " << r.evaluations << "\nStatus: " << r.status << '\n';
}

struct TourResult { std::vector<int> tour; double length; long long evaluations; int iterations; };
// D is symmetric with positive off-diagonal distances. Indices are zero based.
TourResult ant_colony(const Mat& D, Random& rng, int ants, int iterations,
                      double alpha, double beta, double rho, double Q, double tau0, double elite) {
    int n=static_cast<int>(D.size());
    require(n>=3 && ants>0 && iterations>0 && std::isfinite(alpha+beta+Q+tau0+elite)
            && alpha>=0 && beta>=0 && elite>=0 && Q>0 && tau0>0 && rho>0 && rho<1,
            "Invalid ant-colony settings.");
    for (const Vec& row:D) require(row.size()==D.size(), "Distances must form a square matrix.");
    for (int i=0;i<n;++i) for (int j=0;j<n;++j)
        require(std::isfinite(D[i][j]) && D[i][j]==D[j][i] && (i==j ? D[i][j]==0 : D[i][j]>0),
                "Distances must be symmetric, zero on diagonal, and positive otherwise.");
    Mat tau(n,Vec(n,tau0)); TourResult r; r.length=std::numeric_limits<double>::infinity();
    auto deposit=[&](const std::vector<int>& tour,double amount) {
        for (int j=0;j<n;++j) { int a=tour[j],b=tour[(j+1)%n]; tau[a][b]+=amount; tau[b][a]+=amount; }
    };
    for (int k=0;k<iterations;++k) {
        std::vector<std::vector<int>> tours; Vec lengths;
        for (int m=0;m<ants;++m) {
            std::vector<int> tour{rng.index(n)}; std::vector<bool> visited(n,false); visited[tour[0]]=true;
            while (static_cast<int>(tour.size())<n) {
                int current=tour.back(); std::vector<int> available; Vec logw;
                for (int j=0;j<n;++j) if (!visited[j]) {
                    available.push_back(j); logw.push_back(alpha*std::log(tau[current][j])-beta*std::log(D[current][j]));
                }
                double largest=*std::max_element(logw.begin(),logw.end()), total=0;
                Vec weights(logw.size());
                for (std::size_t j=0;j<logw.size();++j) { weights[j]=std::exp(logw[j]-largest); total+=weights[j]; }
                double target=rng.u()*total, cumulative=0; int city=available.back();
                for (std::size_t j=0;j<weights.size();++j) {
                    cumulative+=weights[j]; if (target<cumulative) { city=available[j]; break; }
                }
                tour.push_back(city); visited[city]=true;
            }
            double length=0;
            for (int j=0;j<n;++j) length+=D[tour[j]][tour[(j+1)%n]];
            require(std::isfinite(length) && length>0, "Tour-length overflow; rescale distances.");
            if (length<r.length) { r.tour=tour; r.length=length; }
            tours.push_back(tour); lengths.push_back(length);
        }
        for (Vec& row:tau) for (double& t:row) t*=1-rho;
        for (int m=0;m<ants;++m) deposit(tours[m],Q/lengths[m]);
        deposit(r.tour,elite*Q/r.length);
        for (Vec& row:tau) for (double& t:row) {
            require(std::isfinite(t), "Pheromone overflow; rescale settings.");
            t=std::max(t,std::numeric_limits<double>::min());
        }
    }
    std::rotate(r.tour.begin(),std::find(r.tour.begin(),r.tour.end(),0),r.tour.end());
    r.evaluations=static_cast<long long>(ants)*iterations; r.iterations=iterations; return r;
}

int main() {
  try {
    const unsigned int seed = 7;
    Random rng(seed);
    // Alternative: comment out the previous line and seed from system entropy.
    // Random rng(std::random_device{}());
    Mat xy={{.10,.20},{.85,.15},{.65,.75},{.10,.70},{.55,.70},{.60,.35},
            {.20,.55},{.80,.50},{.42,.12},{.50,.50},{.35,.35},{.40,.60}};
    int n=static_cast<int>(xy.size()); Mat D(n,Vec(n,0));
    for (int i=0;i<n;++i) for (int j=i+1;j<n;++j)
        D[i][j]=D[j][i]=std::hypot(xy[i][0]-xy[j][0],xy[i][1]-xy[j][1]);
    int n_ants=96,max_iter=300; double alpha=1,beta=3,rho=.2,Q=1,tau0=.01,elite_weight=4;
    TourResult r=ant_colony(D,rng,n_ants,max_iter,alpha,beta,rho,Q,tau0,elite_weight);
    std::cout << std::setprecision(12) << "Best length: " << r.length << "\nClosed tour (labels start at 1): ";
    for (int city:r.tour) std::cout << city+1 << " -> ";
    std::cout << r.tour[0]+1 << "\nTours evaluated: " << r.evaluations << "\nIterations: " << r.iterations
              << "\nStatus: Search budget completed.\nRepeat with other seeds; global optimality is not guaranteed.\n";
    return 0;
  } catch (const std::exception& e) {
    std::cerr << "Error: " << e.what() << '\n'; return 1;
  }
}