#!/usr/bin/env python
#
# tsp.py
# ------
#
# Travelling Salesman Problem - use a genetic algorithm to find solutions
#
import random
import re
import sys

import matplotlib.pyplot as plt
import numpy as np
from genoa import GeneticAlgorithm, load_config_parser
from individual import OrderedIndividual


class Point(complex):
    """2-D point"""
    x = property(lambda self: self.real)
    y = property(lambda self: self.imag)

    def __repr__(self):
        return 'Point({}, {})'.format(self.x, self.y)

    def __iter__(self):
        yield self.x
        yield self.y


def distance(a, b):
    """Return the distance between two points (rounded as per TSP datafile)"""
    return int(abs(a - b) + 0.5)


def create_grid(cities):
    """Return a grid of distance between each pair of cities"""
    n = len(cities)
    grid = np.zeros((n, n), dtype=int)
    for i in range(n):
        for j in range(i + 1, n):
            grid[i, j] = grid[j, i] = distance(cities[i], cities[j])
    return grid


def load_tsp_data(filename):
    """Load tsp datafile and return a list of points"""
    pts = []
    tag_found = False

    with open(filename, "r") as fh:
        for line in fh:
            if re.match('EOF', line):
                break

            if re.match('EDGE_WEIGHT_TYPE', line):
                _, t = line.strip().split(': ')
                if t != 'EUC_2D':
                    raise TypeError('TSP format not supported: {}'.format(t))
                continue

            if re.match('NODE_COORD_SECTION', line):
                tag_found = True
                continue

            if not tag_found:
                continue

            # load the coordinates
            _, x, y = line.strip().split()
            pts.append(Point(float(x), float(y)))
    return pts


def tour_length(cities, tour):
    """Return the total length of a given tour"""
    return sum(distance(cities[tour[i - 1]], cities[tour[i]])
               for i in range(len(tour)))


def plot_lines(points, style='bo-'):
    """Plot lines to connect a series of points."""
    plt.plot([p.x for p in points], [p.y for p in points], style)
    # plt.axis('scaled'); plt.axis('off')
    plt.xticks([]);
    plt.yticks([])


def plot_tour(cities, tour):
    """Plot the cities as circles and the tour as lines between them.
    Start city is red square.
    """
    points = [cities[c] for c in tour]
    start = cities[tour[0]]
    plot_lines(points + [start])

    # Mark the start city with a red square
    plot_lines([start], 'rs')


def nearest_neighbour_tour(cities, start=0):
    """Nearest neighbour algorithm - at each point, find the nearest neighbour
    among the unvisited cities and move to it.
    """
    tour = [start]
    remaining = set(range(len(cities)))
    remaining.remove(start)

    while remaining:
        n = min((distance(cities[c], cities[tour[-1]]), c) for c in remaining)
        tour.append(n[1])
        remaining.remove(n[1])
    return tour


def rotate_tour(tour):
    """Rotate a given tour to start at the final city."""
    i = tour.index(max(tour))
    return tour[i:] + tour[:i]


def generate_nntours(cities, howmany):
    """Generate given number of tours starting from random cities."""
    tours = list(rotate_tour(nearest_neighbour_tour(cities, s))
                 for s in random.sample(range(len(cities)), howmany))
    return tours


def create_seed_population(cities, howmany):
    """Create a seed file with tours generated by the nearest-neighbour
    algorithm.
    """
    attr = OrderedIndividual.get_attributes()
    attr['osi.num_genes'] = len(cities) - 1

    tours = generate_nntours(cities, howmany)
    pop = []
    for i in range(len(tours)):
        x = OrderedIndividual(attr)
        # assumes that the tour will start from the final city
        x.genes = tours[i][1:]
        pop.append(x)

    # OrderedIndividual.save_population(filename, pop)
    return pop


if __name__ == '__main__':

    def calculate_fitness(indiv):
        global grid, citylist

        route = [anchor] + indiv.chromosome()

        # segments = [distance(citylist[route[i - 1]], citylist[route[i]])
        #            for i in range(len(route))]
        segments = [grid[route[i - 1], route[i]] for i in range(len(route))]

        indiv.f1 = sum(segments)
        indiv.f2 = min(segments)
        indiv.f3 = max(segments)

        # flip to a maximisation problem
        indiv.fitness = fitness_cap - indiv.f1


    # main
    if len(sys.argv) != 2:
        print('usage: {0} configfile'.format(sys.argv[0]))
        sys.exit(2)
    configfile = sys.argv[1]
    cfp = load_config_parser(configfile)

    citylist = load_tsp_data(cfp.get('Task', 'tsp.datafile'))
    grid = create_grid(citylist)
    anchor = cfp.getint('Task', 'tsp.path_length')
    fitness_cap = cfp.getfloat('Task', 'tsp.cap', fallback=1e6)

    ga = GeneticAlgorithm(cfp, OrderedIndividual, calculate_fitness)
    ga.initialise()
    while ga.iterate():
        pass

    print('\nFinished. GA solution:')
    print(ga.best_individual())
