module tsputils

  !
  ! Various functions for juggling TSP tours.
  ! (evaluation, initialisation, perturbation)
  !

  use tspglobals

  implicit none

  contains

    !
    ! calculate length of given tour
    !
    real function tour_len(tour)
      integer, dimension(0:n_cities-1), intent(in) :: tour

      integer :: i

      tour_len = 0
      do i = 0, n_cities-1
         tour_len = tour_len + calc_dist(tour(i), tour(mod(i+1, n_cities)))
      end do
    end function tour_len

    !
    ! calculate distance between given cities
    !
    real function calc_dist(city1, city2)
      integer, intent(in) :: city1, city2
      
      calc_dist = sqrt( &
           & (city_x(city1) - city_x(city2))**2 + &
           & (city_y(city1) - city_y(city2))**2 )
    end function calc_dist

    !
    ! verify that a tour is valid
    !
    logical function validate_tour(tour)
      integer, dimension(0:n_cities-1), intent(in) :: tour
      integer :: i, j
      required_cities: do i = 0, n_cities - 1
         search_for_city: do j = 0, n_cities - 1
            if (tour(j) == i) then
               ! found city i
               cycle required_cities
            end if
         end do search_for_city
         ! failed to find city i
         validate_tour = .false.
         return
      end do required_cities
      ! all required cities found
      validate_tour = .true.
    end function validate_tour

    !
    ! generate a tour in numerical order
    !
    subroutine numerical_tour(tour)
      integer, dimension(0:n_cities-1), intent(out) :: tour
      integer :: i
      tour = (/ (i,i=0,n_cities-1) /)
    end subroutine numerical_tour

    !
    ! generate a random tour
    !
    subroutine random_tour(tour)
      integer, dimension(0:n_cities-1), intent(out) :: tour
      real :: rand
      integer :: i, pool, next
      call numerical_tour(city_order_temp)
      pool = n_cities ! remaining cities to be ordered

      do i = 0, n_cities - 1
         ! select random city from those remaining
         call random_number(rand)
         next = rand * (pool) ! index in pool of next city
         tour(i) = city_order_temp(next)
         ! shuffle up so remaining cities are in city_order_temp(0:pool-1)
         pool = pool - 1
         city_order_temp(next:pool-1) = city_order_temp(next+1:pool)
      end do
    end subroutine random_tour

    !
    ! generate a greedy tour from random starting point
    !
    subroutine greedy_tour(tour)
      integer, dimension(0:n_cities-1), intent(out) :: tour
      real :: rand
      integer :: i, pool, next, temp
      integer :: j
      real :: dist, next_dist

      call numerical_tour(city_order_temp)
      pool = n_cities ! remaining cities to be ordered

      ! choose starting point at random
      call random_number(rand)
      next = rand * (n_cities)
      call swap_operator(city_order_temp, next, n_cities-1)
      tour(0) = city_order_temp(n_cities-1)
      pool = pool - 1

      do i = 1, n_cities - 1
         ! select closest city from those remaining
         next_dist = huge(next_dist)
         do j = 0, pool - 1
            dist = calc_dist(tour(i-1), city_order_temp(j))
            if (dist < next_dist) then
               next_dist = dist
               next = j ! index in pool of next city
            end if
         end do
         tour(i) = city_order_temp(next)
         ! shuffle up so remaining cities are in city_order_temp(0:pool-1)
         pool = pool - 1
         city_order_temp(next:pool-1) = city_order_temp(next+1:pool)
      end do
    end subroutine greedy_tour

    !
    ! swap cities at given cities in given tour
    !
    subroutine swap_operator(tour, i1, i2)
      integer, dimension(0:n_cities-1), intent(inout) :: tour
      integer, intent(in) :: i1, i2

      integer :: temp

      ! swap these 2 cities
      temp = tour(i1)
      tour(i1) = tour(i2)
      tour(i2) = temp
    end subroutine swap_operator

    !
    ! calculate resulting distance from swapping cities at given indices
    !
    real function swapped_distance(tour, distance, i1, i2)
      integer, dimension(0:n_cities-1), intent(in) :: tour
      real, intent(in) :: distance
      integer, intent(in) :: i1, i2
      
      integer :: i1_pre, i1_post, i2_pre, i2_post

      ! catch swap with self
      if (i1 == i2) then
         swapped_distance = distance
         return
      end if

      i1_pre =  mod(( i1 - 1 )+n_cities, n_cities)
      i1_post = mod(( i1 + 1 ),          n_cities)
      i2_pre =  mod(( i2 - 1 )+n_cities, n_cities)
      i2_post = mod(( i2 + 1 ),          n_cities)

      swapped_distance = distance

      ! cut existing edges
      swapped_distance = swapped_distance &
           & - calc_dist(tour(i1), tour(i1_pre)) &
           & - calc_dist(tour(i1), tour(i1_post)) &
           & - calc_dist(tour(i2), tour(i2_pre)) &
           & - calc_dist(tour(i2), tour(i2_post))

      ! add in new edges
      swapped_distance = swapped_distance &
           & + calc_dist(tour(i2), tour(i1_pre)) &
           & + calc_dist(tour(i2), tour(i1_post)) &
           & + calc_dist(tour(i1), tour(i2_pre)) &
           & + calc_dist(tour(i1), tour(i2_post))

      ! if they are adjacent then the swapping modifies this
      if ((abs(i2 - i1) == 1) .or. (abs(i2 - i1) == n_cities - 1)) then
         swapped_distance = swapped_distance &
              & + 2 * calc_dist(tour(i1), tour(i2))
      end if

    end function swapped_distance

    !
    ! reverse segment of tour between given cities [i1, i2)
    !
    subroutine reverse_operator(tour, i1, i2)
      integer, dimension(0:n_cities-1), intent(inout) :: tour
      integer, intent(in) :: i1, i2

      integer :: seg_len, k, k1, k2, temp

      ! calculate segment length
      if (i2 > i1) then
         seg_len = i2 - i1
      else
         seg_len = i2 - i1 + n_cities
      end if

      ! no point in reversing whole path
      if (seg_len >= n_cities - 2) then
         return
      end if

      ! swap elements starting from extremes of segment
      swap_loop: do k = 0, seg_len / 2 - 1
         k1 = i1 + k
         k2 = i2 - 1 - k
         ! wrap-around the indices to swap
         k1 = mod(k1 + n_cities, n_cities)
         k2 = mod(k2 + n_cities, n_cities)
         ! swap them
         temp = tour(k1)
         tour(k1) = tour(k2)
         tour(k2) = temp
      end do swap_loop

    end subroutine reverse_operator

    !
    ! calculate resulting distance from reversing segment of tour: [i1, i2)
    !
    real function reversed_distance(tour, distance, i1, i2)
      integer, dimension(0:n_cities-1), intent(in) :: tour
      real, intent(in) :: distance
      integer, intent(in) :: i1, i2
      
      integer :: seg_len, i1_pre, i2_pre

      ! calculate segment length
      if (i2 > i1) then
         seg_len = i2 - i1
      else
         seg_len = i2 - i1 + n_cities
      end if

      ! no point in reversing whole path
      if (seg_len >= n_cities - 2) then
         reversed_distance = distance
         return
      end if

      i1_pre = mod(( i1 - 1 )+n_cities, n_cities)
      i2_pre = mod(( i2 - 1 )+n_cities, n_cities)

      reversed_distance = distance

      !print *, 'pre rev', reversed_distance

      ! cut existing edges
      reversed_distance = reversed_distance &
           & - calc_dist(tour(i1), tour(i1_pre)) &
           & - calc_dist(tour(i2_pre), tour(i2))

      !print *, tour(i1), tour(i1_pre), calc_dist(tour(i1), tour(i1_pre))
      !print *, tour(i2_pre), tour(i2), calc_dist(tour(i2_pre), tour(i2))
      !print *, 'rm old', reversed_distance

      ! add in new edges
      reversed_distance = reversed_distance &
           & + calc_dist(tour(i2_pre), tour(i1_pre)) &
           & + calc_dist(tour(i1), tour(i2))

      !print *, tour(i2_pre), tour(i1_pre), calc_dist(tour(i2_pre), tour(i1_pre))
      !print *, tour(i1), tour(i2), calc_dist(tour(i1), tour(i2))
      !print *, 'add new', reversed_distance

    end function reversed_distance

    !
    ! shift segment of tour [i1, i2) to after i3
    !
    subroutine shift_operator(tour, i1, i2, i3)
      integer, dimension(0:n_cities-1), intent(inout) :: tour
      integer, intent(in) :: i1, i2, i3

      integer :: seg_len, ex_seg_len, k, k_src, k_dst

      ! calculate segment length
      if (i2 > i1) then
         seg_len = i2 - i1
      else
         seg_len = i2 - i1 + n_cities
      end if

      ! no point in shifting whole path
      if (seg_len >= n_cities - 2) then
         return
      end if

      city_order_temp = tour

      ! shift segment
      do k = 0, seg_len-1
         k_src = mod(( i1 + k ), n_cities)
         k_dst = mod(( i3 + k + 1 ), n_cities)
         tour(k_dst) = city_order_temp(k_src)
      end do

      ! realign region between previous segment and destination
      if (i1 > i3) then
         ex_seg_len = i1 - i3 - 1
      else
         ex_seg_len = i1 - i3 - 1 + n_cities
      end if

      do k = 0, ex_seg_len-1
         k_src = mod(( i3 + k + 1 ), n_cities)
         k_dst = mod(( i3 + k + 1 + seg_len ), n_cities)
         tour(k_dst) = city_order_temp(k_src)
      end do
    end subroutine shift_operator

    !
    ! calculate resulting distance from shifting segment of tour: 
    ! [i1, i2) to after i3 (which is randomly chosen and returned)
    !
    real function shifted_distance(tour, distance, i1, i2, i3)
      integer, dimension(0:n_cities-1), intent(in) :: tour
      real, intent(in) :: distance
      integer, intent(in) :: i1, i2
      integer, intent(out) :: i3
      
      integer :: seg_len, i1_pre, i2_pre, i3_post
      real :: rand

      ! calculate segment length
      if (i2 > i1) then
         seg_len = i2 - i1
      else
         seg_len = i2 - i1 + n_cities
      end if

      ! no point in shifting whole path
      if (seg_len >= n_cities - 2) then
         i3 = 0
         shifted_distance = distance
         return
      end if

      ! choose an index i3 to shift to (after)
      CALL random_number(rand)
      ! random offset
      i3 = rand * (n_cities - seg_len - 1)
      i3 = mod(( i2 + i3 ), n_cities)

      shifted_distance = distance

      i1_pre = mod(( i1 - 1 )+n_cities, n_cities)
      i2_pre = mod(( i2 - 1 )+n_cities, n_cities)
      i3_post = mod(( i3 + 1 )+n_cities, n_cities)

      ! cut existing edges
      shifted_distance = shifted_distance &
           & - calc_dist(tour(i1), tour(i1_pre)) &
           & - calc_dist(tour(i2_pre), tour(i2))
      ! cut edge at shift destination
      shifted_distance = shifted_distance &
           & - calc_dist(tour(i3), tour(i3_post))

      ! splice in segment
      shifted_distance = shifted_distance &
           & + calc_dist(tour(i1), tour(i3)) &
           & + calc_dist(tour(i2_pre), tour(i3_post))
      ! close hole where segment came from
      shifted_distance = shifted_distance &
           & + calc_dist(tour(i1_pre), tour(i2))
    end function shifted_distance

    !
    ! Metropolis algorithm for the Boltzmann distribution
    !
    logical function metrop_boltz(cost, temperature)
      real, intent(in) :: cost, temperature
      real :: rand
      if (cost <= 0) then
         metrop_boltz = .true.
      else if (temperature == 0) then
         metrop_boltz = .false.
      else
         call random_number(rand)
         metrop_boltz = (rand < exp( - cost/temperature ))
      end if
    end function metrop_boltz

end module tsputils