! ch13_ant_colony.f95
! 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. Edit coordinates and
! search settings in the driver. Distances use the coordinate units.
! Repeat with different seeds; global optimality is not guaranteed.
! Console output only. Uses double precision and Fortran 95 syntax.
!
! Compile:
! gfortran -std=f95 -Wall -Wextra -O2 ch13_ant_colony.f95 -o ch13_ant_colony
! Run: ./ch13_ant_colony (Windows: ch13_ant_colony.exe)

module ant_colony_methods
    implicit none
    private
    integer, parameter :: dp=selected_real_kind(15,307)
    public :: dp, initialize_random, ant_colony
contains
    subroutine initialize_random(seed)
        integer, intent(in) :: seed
        integer, allocatable :: seeds(:)
        integer :: n,i
        call random_seed(size=n)
        allocate(seeds(n))
        do i=1,n
            seeds(i)=1+modulo(modulo(seed,1000000)+37*modulo(i,1000000),1000000)
        end do
        call random_seed(put=seeds)
        deallocate(seeds)
        ! Repeatable on the same compiler/runtime; sequences depend on compiler.
    end subroutine initialize_random

    subroutine ant_colony(D,M,K,alpha,beta,rho,Q,tau0,elite_weight,best_tour,best_length)
        real(dp), intent(in) :: D(:,:),alpha,beta,rho,Q,tau0,elite_weight
        integer, intent(in) :: M,K
        integer, intent(out) :: best_tour(:)
        real(dp), intent(out) :: best_length
        real(dp) :: tau(size(D,1),size(D,1)),lengths(M)
        real(dp) :: logw(size(D,1)),weights(size(D,1)),u,total,cumulative
        integer :: tours(M,size(D,1)),available(size(D,1)),rotated(size(D,1))
        integer :: n,i,j,ant,iteration,step,current,nav,choice,idx(1),first
        logical :: visited(size(D,1))

        n=size(D,1)
        if (n<3 .or. size(D,2)/=n .or. size(best_tour)/=n) stop 'Invalid distance-matrix dimensions.'
        if (.not. all(finite_value(D))) stop 'Distances must be finite.'
        do i=1,n
            if (D(i,i)/=0.0_dp) stop 'Distance matrix must have zero diagonal.'
            do j=i+1,n
                if (D(i,j)<=0.0_dp .or. D(i,j)/=D(j,i)) stop 'Require positive symmetric edge distances.'
            end do
        end do
        if (M<1 .or. K<1) stop 'Ant and iteration counts must be positive.'
        if (.not. finite_value(alpha) .or. alpha<0.0_dp) stop 'Invalid pheromone exponent.'
        if (.not. finite_value(beta) .or. beta<0.0_dp) stop 'Invalid distance exponent.'
        if (.not. finite_value(rho) .or. rho<=0.0_dp .or. rho>=1.0_dp) stop 'Invalid evaporation rate.'
        if (.not. finite_value(Q) .or. Q<=0.0_dp) stop 'Invalid deposit scale.'
        if (.not. finite_value(tau0) .or. tau0<=0.0_dp) stop 'Invalid initial pheromone.'
        if (.not. finite_value(elite_weight) .or. elite_weight<0.0_dp) stop 'Invalid elite weight.'

        tau=tau0
        best_length=huge(1.0_dp)
        best_tour=0
        do iteration=1,K
            do ant=1,M
                visited=.false.
                call random_number(u)
                tours(ant,1)=1+int(n*u)
                visited(tours(ant,1))=.true.
                do step=2,n
                    current=tours(ant,step-1)
                    nav=0
                    do j=1,n
                        if (.not. visited(j)) then
                            nav=nav+1
                            available(nav)=j
                            logw(nav)=alpha*log(tau(current,j))-beta*log(D(current,j))
                        end if
                    end do
                    ! Relative weights proportional to pheromone^alpha / distance^beta.
                    ! Subtracting the largest log weight avoids underflow of all weights.
                    if (.not. all(finite_value(logw(1:nav)))) stop 'Rescale pheromone or distance settings.'
                    weights(1:nav)=exp(logw(1:nav)-maxval(logw(1:nav)))
                    total=sum(weights(1:nav))
                    call random_number(u)
                    cumulative=0.0_dp
                    choice=nav  ! Rounding-safe fallback to the last available city.
                    do j=1,nav
                        cumulative=cumulative+weights(j)/total
                        if (u<cumulative) then
                            choice=j
                            exit
                        end if
                    end do
                    tours(ant,step)=available(choice)
                    visited(available(choice))=.true.
                end do
                lengths(ant)=0.0_dp
                do step=1,n
                    i=tours(ant,step)
                    j=tours(ant,mod(step,n)+1)
                    lengths(ant)=lengths(ant)+D(i,j)
                end do
                if (.not. finite_value(lengths(ant))) stop 'Route length overflow: rescale distances.'
            end do
            idx=minloc(lengths)
            if (lengths(idx(1))<best_length) then
                best_length=lengths(idx(1))
                best_tour=tours(idx(1),:)
            end if
            ! Complete all routes before evaporating and depositing pheromone.
            tau=(1.0_dp-rho)*tau
            do ant=1,M
                call deposit(tau,tours(ant,:),Q/lengths(ant))
            end do
            call deposit(tau,best_tour,elite_weight*Q/best_length)
            if (.not. all(finite_value(tau))) stop 'Pheromone overflow: rescale deposit settings.'
            tau=max(tau,tiny(1.0_dp))
        end do
        ! Rotate the display order so city 1 comes first; length is unchanged.
        first=1
        do i=1,n
            if (best_tour(i)==1) first=i
        end do
        do i=1,n
            rotated(i)=best_tour(mod(first+i-2,n)+1)
        end do
        best_tour=rotated
    end subroutine ant_colony

    subroutine deposit(tau,route,amount)
        real(dp), intent(inout) :: tau(:,:)
        integer, intent(in) :: route(:)
        real(dp), intent(in) :: amount
        integer :: s,i,j,n
        n=size(route)
        do 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 do
    end subroutine deposit

    elemental logical function finite_value(x)
        real(dp), intent(in) :: x
        finite_value=abs(x)<=huge(x)
    end function finite_value
end module ant_colony_methods

program ch13_ant_colony
    use ant_colony_methods, only: dp, initialize_random, ant_colony
    implicit none
    ! Edit city count and coordinates together. Each column is one city (x,y).
    integer, parameter :: n=12,n_ants=96,max_iter=300,seed=7
    real(dp), parameter :: XY(2,n)=reshape((/ &
        0.10_dp,0.20_dp, 0.85_dp,0.15_dp, 0.65_dp,0.75_dp, &
        0.10_dp,0.70_dp, 0.55_dp,0.70_dp, 0.60_dp,0.35_dp, &
        0.20_dp,0.55_dp, 0.80_dp,0.50_dp, 0.42_dp,0.12_dp, &
        0.50_dp,0.50_dp, 0.35_dp,0.35_dp, 0.40_dp,0.60_dp /),(/2,n/))
    real(dp), parameter :: alpha=1.0_dp,beta=3.0_dp,rho=0.20_dp
    real(dp), parameter :: Q=1.0_dp,tau0=0.01_dp,elite_weight=4.0_dp
    real(dp) :: D(n,n),best_length,delta(2),scale
    integer :: tour(n),i,j
    ! integer :: clock_seed  ! Uncomment for clock-based seeding.

    call initialize_random(seed)
    ! Alternative: comment out the line above, uncomment the clock_seed
    ! declaration above, and uncomment both lines below.
    ! call system_clock(count=clock_seed)
    ! call initialize_random(clock_seed)

    D=0.0_dp
    do i=1,n
        do j=i+1,n
            delta=XY(:,i)-XY(:,j)
            scale=maxval(abs(delta))
            if (scale==0.0_dp) stop 'Cities must have distinct coordinates.'
            D(i,j)=scale*sqrt(sum((delta/scale)**2))
            D(j,i)=D(i,j)
        end do
    end do
    call ant_colony(D,n_ants,max_iter,alpha,beta,rho,Q,tau0,elite_weight,tour,best_length)
    write(*,'(/,a)') 'Ant colony optimization'
    write(*,'(a,i0)') 'Cities:            ',n
    write(*,'(a,es23.15)') 'Best route length: ',best_length
    write(*,'(a,i0)') 'Iterations:        ',max_iter
    write(*,'(a,i0)') 'Routes evaluated:  ',n_ants*max_iter
    write(*,'(a)') 'Best closed route (starting at city 1):'
    do i=1,n
        write(*,'(i0,1x)',advance='no') tour(i)
    end do
    write(*,'(i0)') tour(1)
    write(*,'(a)') 'Status: Search budget completed.'
    write(*,'(a)') 'Repeat with other seeds; global optimality is not guaranteed.'
end program ch13_ant_colony