"""
atsp.py: Construction and local optimization for the ATSP.

The Asymmetric Traveling Salesman Problem (ATSP) is a combinatorial
optimization problem, where given a distance matrix, one wants to find
an order for visiting all the cities in such a way that the travel
distance is minimal.

This file contains a set of functions to illustrate:
  - construction heuristics for the ATSP
  - improvement heuristics for a previously constructed solution
  - local search, and random-start local search.

Copyright (c) by Joao Pedro PEDROSO, 2010
"""
from __future__ import with_statement

import math
import random
import gzip

def read_atsplib(filename):
    "basic function for reading a ATSP problem on the TSPLIB format"
    "NOTE: only works for explicit matrices"

    if filename[-3:] == ".gz":
        f = gzip.open(filename, 'r')
        data = f.readlines()
    else:
        with open(filename, 'r') as f:
            data = f.readlines()

    for line in data:
        if line.find("DIMENSION") >= 0:
            n = int(line.split()[1])
            # print "dimension:", n
            break
    else:
        raise IOError("'DIMENSION' keyword not found in file '%s'" % filename)

    for line in data:
        if line.find("EDGE_WEIGHT_TYPE") >= 0:
            if line.split()[1] == "EXPLICIT":
                # print "explicit"
                break
    else:
        raise IOError("'EDGE_WEIGHT_TYPE' is not 'EXPLICIT' in file '%s'" % filename)
    
    for k,line in enumerate(data):
        if line.find("EDGE_WEIGHT_SECTION") >= 0:
            # print "found weights"
            break
    else:
        raise IOError("'EDGE_WEIGHT_SECTION' not found in file '%s'" % filename)

    D = {}      # dictionary to hold n times n matrix
    # flatten list of distances
    dist = []
    for line in data[k+1:]:
        if line.find("EOF") >= 0:
            break
        for val in line.split():
            dist.append(int(val))

    k = 0
    for i in range(n):
        for j in range(n):
            D[i,j] = dist[k]
            k += 1

    return n, D


def mk_closest(D, n, m):
    """Compute a sorted list of the 'm' closest distances for each of the nodes.

    For each node, the entry is in the form [(d1,i1), (d2,i2), ...]
    where each tuple is a pair (distance,node).

    returns (I,O)  where 'I' is for in nodes, 'O' for out nodes.
    """
    I = []	# lowest in-arcs
    for i in range(n):
        dlist = [(D[j,i], j) for j in range(n) if j != i]
        dlist.sort()
        I.append(dlist[:m])
    O = []	# lowest out-arcs
    for i in range(n):
        dlist = [(D[i,j], j) for j in range(n) if j != i]
        dlist.sort()
        O.append(dlist[:m])
    return I,O


def length(sol, D):
    """Calculate the length of a tour according to distance matrix 'D'."""
    z = D[sol[-1], sol[0]]    # edge from last to first city of the sol
    for i in range(1,len(sol)):
        z += D[sol[i-1], sol[i]]      # add length of edge from city i-1 to i
    return z


def randtour(unvisited):
    """Construct a random sol of size 'n'."""
    sol = list(unvisited)
    random.shuffle(sol) # place it in a random order
    return sol


def nearest(unvisited, last, D):
    """Return the index of the node which is closest to 'last'."""
    near = unvisited[0]
    min_dist = D[last, near]
    for i in unvisited[1:]:
        if D[last,i] < min_dist:
            near = i
            min_dist = D[last, near]
    return near


def nearest_neighbor(unvisited_, i, D):
    """Return sol starting from city 'i', going through all 'unvisited', using the Nearest Neighbor.

    Uses the Nearest Neighbor heuristic to construct a solution:
    - start visiting city i
    - while there are unvisited cities, follow to the closest one
    - return to city i
    """
    unvisited = list(unvisited_)
    unvisited.remove(i)
    last = i
    sol = [i]
    while unvisited != []:
        next = nearest(unvisited, last, D)
        sol.append(next)
        unvisited.remove(next)
        last = next
    return sol


def greedy(nodes, D):
    """Return sol constructed with the Greedy heuristics.

    Uses the greedy heuristics to construct a solution:
    - add allowed nodes until all nodes are connected
    """
    n = len(nodes)
    arcs = [(D[i,j],i,j) for i in nodes for j in nodes if j != i]
    arcs.sort()
    from_set = set([])
    to_set = set([])
    dist, narcs = 0,0
    succ = [None for i in nodes]
    while arcs != []:
        d,i,j = arcs.pop(0)
        if i in from_set or j in to_set:
            continue
        allowed = True
        if narcs < n-1:
            j_ = succ[j]
            while j_ != None:
                if i == j_:
                    allowed = False
                    break
                j_ = succ[j_]
        if not allowed:
            continue
        succ[i] = j
        from_set.add(i)
        to_set.add(j)
        dist += d
        narcs += 1
        print "\tadded", (i,j), "--->", dist
    i = 0
    sol = [i]
    j = succ[i]
    while j != 0:
        sol.append(j)
        j = succ[j]
    assert len(sol) == n
    return sol, dist


def improve_simplistic(sol, z, D):
    """Try to improve 'sol' by exchanging arcs; return improved sol length.
    
    If possible, make a series of local improvements on the solution 'sol',
    using a first-improve strategy, until reaching a local optimum.
    """
    n = len(sol)
    for i in range(n-2):
        a,b = sol[i],sol[i+1]
        for j in range(i+1,n-1):
            c,d = sol[j],sol[j+1]
            for k in range(j+1,n):
                e,f = sol[k],sol[(k+1)%n]
                delta = (D[a,d] + D[c,f] + D[e,b]) - (D[a,b] + D[c,d] + D[e,f])
                if delta < 0:
                    sol = sol[0:i+1] + sol[j+1:k+1] + sol[i+1:j+1] + sol[k+1:] 
                    a,b = sol[i],sol[i+1]   # must update after exchange
                    c,d = sol[j],sol[j+1]
                    z += delta
                    # print "\t * ", sol, z
    return sol, z


def improve(sol, tinv, z, D, I, O):
    """Try to improve 'sol' by exchanging arcs; return improved sol length.
    
    If possible, make a series of local improvements on the solution 'sol',
    using a best first strategy, until reaching a local optimum.

    Sol is represented as (->-ab-->--cd-->--ef-->--), and the only possible move
    is a 3-exchange into (->-ad-->--eb-->--cf-->--), for having to path inversions.    
    """
    n = len(sol)
    N = range(n); random.shuffle(N)
    for j in N:		# j is the index of the middle arc
        c, d = sol[j-1], sol[j]
        dist_cd = D[c,d]
        improved = False

        for dist_cf,f in O[c]:
            if improved or dist_cf >= dist_cd:
                break
            k = tinv[f]; assert sol[k] == f
            assert k != j-1
            e = sol[k-1]
            dist_ef = D[e,f]

            for dist_eb,b in O[e]:
                if improved or dist_eb + dist_cf >= dist_cd + dist_ef:
                    break
                i = tinv[b]; assert sol[i] == b
                if i <= k <= j or k <= j <= i or j <= i <= k:   # infeasible
                    continue
                a = sol[i-1]
                dist_ab = D[a,b]
                dist_ad = D[a,d]

                delta = (dist_ad + dist_eb + dist_cf) - (dist_ab + dist_cd + dist_ef)
                if delta < 0:       # exchange decreases length
                    idx = sorted([i,j,k])
                    ii,jj,kk = tuple(idx)
                    sol = sol[0:ii] + sol[jj:kk] + sol[ii:jj] + sol[kk:]
                    # print z, "--->",
                    z += delta
                    # print z
                    for r in range(ii,kk):
                        tinv[sol[r]] = r  # update position of each city in 'sol'
                    improved = True # values for a, b, ... are no long uptodate

        for dist_ad,a in I[d]:
            if improved or dist_ad >= dist_cd:
                break
            i = (tinv[a]+1)%n; assert sol[i-1] == a
            assert i != j
            b = sol[i]
            dist_ab = D[a,b]

            for dist_eb,e in I[b]:
                if improved or dist_eb + dist_ad >= dist_ab + dist_cd:
                    break
                k = (tinv[e]+1)%n; assert sol[k-1] == e
                if i <= k <= j or k <= j <= i or j <= i <= k:   # infeasible
                    continue
                f = sol[k]
                dist_ef = D[e,f]
                dist_cf = D[c,f]

                delta = (dist_ad + dist_eb + dist_cf) - (dist_ab + dist_cd + dist_ef)
                if delta < 0:       # exchange decreases length
                    idx = sorted([i,j,k])
                    ii,jj,kk = tuple(idx)
                    sol = sol[0:ii] + sol[jj:kk] + sol[ii:jj] + sol[kk:] 
                    z += delta
                    for r in range(ii,kk):
                        tinv[sol[r]] = r  # update position of each city in 'sol'
                    improved = True # values for a, b, ... are no long uptodate

    return sol, z


def localsearch(sol, z, D, I=None, O=None):
    """Obtain a local optimum starting from solution t; return solution length.

    Parameters:
      sol -- initial sol
      z -- length of the initial sol
      D -- distance matrix
    """
    n = len(sol)
    tinv = [0 for i in sol]
    for k in range(n):
        tinv[sol[k]] = k  # position of each city in 'sol'
    if O == None:
        I,O = mk_closest(D, n, n)     # create a sorted list of distances to each node
    while 1:
        sol, newz = improve(sol, tinv, z, D, I, O)
        if newz < z:
            z = newz
        else:
            break
    assert newz == length(sol,D)
    return sol, z


def multistart_localsearch(k, unvisited, D, report=None):
    """Do k iterations of random-start local search, starting from random solutions.

    Parameters:
    -k -- number of iterations
    -D -- distance matrix
    -report -- if not None, call it to print verbose output

    Returns best solution and its cost.
    """
    I,O = mk_closest(D, n, n) # create a sorted list of distances to each node
    bestt=None
    bestz=None
    for i in range(0,k):
        sol = randtour(unvisited)
        z = length(sol, D)
        sol, z = localsearch(sol, z, D, I, O)
        if z < bestz or bestz == None:
            bestz = z
            bestt = list(sol)
            if report:
                report(z, sol)

    return bestt, bestz


def iterated_localsearch(k, unvisited, D, report=None):
    """Do k iterations of iterated local search.

    Parameters:
    -k -- number of iterations
    -D -- distance matrix
    -report -- if not None, call it to print verbose output

    Returns best solution and its cost.
    """
    bestt=None
    bestz=None
    sol = randtour(unvisited)
    n = len(sol)
    z = length(sol, D)
    for it in range(0,k):
        sol, z = localsearch(sol, z, D)
        if z < bestz or bestz == None:
            bestz = z
            bestt = list(sol)
            if report:
                report(z, sol)
        print z, sol[:25]
        # 4-exchange
        R = range(1,n)
        idx = []
        for r in range(4):
            i = random.choice(R); R.remove(i); idx.append(i)
        idx.sort()
        i,j,k,l = tuple(idx)
        a,b,c,d,e,f,g,h = sol[i-1], sol[i], sol[j-1], sol[j], sol[k-1], sol[k], sol[l-1], sol[l]
        delta = D[a,f] + D[g,d] + D[e,b] + D[c,h] - (
            D[a,b] + D[c,d] + D[e,f] + D[g,h])
        z += delta
        sol = sol[:i] + sol[k:l] + sol[j:k] + sol[i:j] + sol[l:]
        print "***", z, "---->",
        # assert z == length(sol, D)
    return bestt, bestz



if __name__ == "__main__":
    """Local search for the Asymmetric Travelling Saleman Problem: sample usage."""
    
    #
    # test the functions:
    #

    import sys
    if len(sys.argv) == 1:
        random.seed(14)    # uncomment for having always the same behavior
        instance = "toy.atsp"
        instance = "INSTANCES/br17.atsp.gz"
        instance = "INSTANCES/p43.atsp.gz"
        # instance = "INSTANCES/kro124p.atsp.gz"
    elif len(sys.argv) == 3:
        seed = int(sys.argv[1])
        random.seed(seed)
        print "seed:", seed
        instance = sys.argv[2]
    else:
        print "usage:", sys.argv[0], "[seed instance maxcpu method]"
    n, D = read_atsplib(instance)
    
    # function for printing best found solution when it is found
    from time import clock
    init = clock()
    def report_sol(obj, s=""):
        print "\ncpu:%r\tobj:%r\tsol:%s" % \
              (clock(), obj, s)


    # print "*** asymmetric travelling salesman problem ***"
    # print
    # print "*** D ***"
    # for i in range(n):
    #     for j in range(n):
    #         print "%5d" % D[i,j],
    #     print
    # print


    # random construction
    print "random construction + local search:"
    unvisited = range(n)
    sol = randtour(unvisited)     # create a random sol
    z = length(sol, D)     # calculate its length
    print "random:", sol, z, '  -->  ',
    sol, z = localsearch(sol, z, D)      # local search starting from the random sol
    print sol, z
    print

    # nearest neighbor construction
    print "construction with nearest neighbor + local search:"
    for i in range(n):
        sol = nearest_neighbor(unvisited, i, D)     # create a NN sol, visiting city 'i' first
        z = length(sol, D)
        print "nneigh:", sol, z, '  -->  ',
        sol, z = localsearch(sol, z, D)
        print sol, z
    print

    # greedy construction
    print "greedy construction + local search:"
    sol, z = greedy(unvisited,D)
    print "greedy:", sol, z, '  -->  ',
    sol, z = localsearch(sol, z, D)      # local search starting from the random sol
    print sol, z

    # multi-start local search
    print "random start local search:"
    niter = 100
    sol,z = multistart_localsearch(niter, unvisited, D, report_sol)
    assert z == length(sol, D)
    print "best found solution (%d iterations): z = %g" % (niter, z)
    print sol
    print

    # iterated local search
    print "iterated local search:"
    niter = 1000
    sol, z = iterated_localsearch(niter, unvisited, D, report_sol)
    print sol, z
    assert z == length(sol, D)

    # check if the code works for a subset of cities to visit
    unvisited = range(n/2)
    z = length(sol, D)
    print "subset local search:"
    sol = randtour(unvisited)    # create a random sol starting from 0
    z = length(sol, D)     # calculate its length
    print "random:", sol, z, '  -->  ',   
    sol, z = localsearch(sol, z, D)      # local search starting from the random sol
    print sol, z
    print
    assert z == length(sol, D)

    print "subset search: greedy construction with nearest neighbor + local search:"
    sol = nearest_neighbor(unvisited, 0, D)     # create a greedy sol, visiting city '0' first
    z = length(sol, D)
    print "nneigh:", sol, z, '  -->  ',
    sol, z = localsearch(sol, z, D)
    print sol, z
    print
    assert z == length(sol, D)