import cPickle as pickle
import heapq
from itertools import combinations, izip
import os

import numpy as np
from matplotlib import pyplot as plt
from matplotlib import image as mpimg

import networkx as nx
import pytsp

os.environ['LKH'] = '/Users/jgoldsmith/Downloads/LKH-2.0.7/LKH'

# Read in the map as an image
img = mpimg.imread("map.png")
radius = 20

def get_bresenham_line(x1, y1, x2, y2, square_is_filled):
  # Setup initial conditions
  dx = x2 - x1
  dy = y2 - y1

  # Determine how steep the line is
  is_steep = abs(dy) > abs(dx)

  # Rotate line
  if is_steep:
    x1, y1 = y1, x1
    x2, y2 = y2, x2

  # Swap start and end points if necessary and store swap state
  swapped = False
  if x1 > x2:
    x1, x2 = x2, x1
    y1, y2 = y2, y1
    swapped = True

  # Recalculate differentials
  dx = x2 - x1
  dy = y2 - y1

  # Calculate error
  error = int(dx / 2.0)
  ystep = 1 if y1 < y2 else -1

  # Iterate over bounding box generating points between start and end
  y = y1
  points = []
  for x in np.arange(x1, x2 + 1):
    coord = (y, x) if is_steep else (x, y)
    if square_is_filled(coord):
      break

    points.append(coord)
    error -= abs(dy)
    if error < 0:
      y += ystep
      error += dx

  # Reverse the list if the coordinates were swapped
  if swapped:
    points.reverse()
  return points

def point_is_filled(point):
  x, y = point

  if 0 <= x < img.shape[1] and 0 <= y < img.shape[0]:
    return not img[y][x].any()  # 0 is black, 1 is white

  return True

def visible_points_from_point(point):
  x, y = point
  points = []

  theta = 0
  for degrees in np.arange(0, 360, 5):
    theta = degrees * np.pi / 180.0

    x1 = radius * np.cos(theta) + x
    y1 = radius * np.sin(theta) + y

    line = get_bresenham_line(x, y, x1, y1, point_is_filled)
    points.extend(line)

  return points

height, width, _ = img.shape
start = (int(width / 2), int(height / 2))

class Map(object):
  def __init__(self, img, tile_size=5):
    self.img = img
    self.tile_size = tile_size

  @property
  def tile_width(self):
    return int(self.img.shape[1] / self.tile_size)

  @property
  def tile_height(self):
    return int(self.img.shape[0] / self.tile_size)

  def tiles(self):
    for x in np.arange(self.tile_height):
      for y in np.arange(self.tile_width):
        yield (x, y)

  def tile_for_point(self, point):
    x, y = point
    tilex = int(np.floor(x / float(self.tile_size)))
    tiley = int(np.floor(y / float(self.tile_size)))

    return (tilex, tiley)

  def points_in_tile(self, tile):
    x, y = tile
    for pointx in np.arange(x * self.tile_size, (x + 1) * self.tile_size):
      for pointy in np.arange(y * self.tile_size, (y + 1) * self.tile_size):
        if 0 <= pointx < self.img.shape[0] and 0 <= pointy < self.img.shape[1]:
          yield (pointx, pointy)

  def point_for_tile(self, tile):
    x, y = tile

    pointx = int((x + 0.5) * self.tile_size)
    pointy = int((y + 0.5) * self.tile_size)

    return (pointx, pointy)

  def point_is_filled(self, point):
    return not self.img[point[0]][point[1]].all()

  def tile_is_filled(self, tile):
    return any(
      self.point_is_filled(point)
      for point in self.points_in_tile(tile)
    )

  def neighbors_for_point(self, point):
    directions = [(0, 1),
                  (0, -1),
                  (1, 0),
                  (-1, 0)] #,
                  # (1, 1),    (uncomment for neighbors8 instead of neighbors4)
                  # (1, -1),
                  # (-1, 1),
                  # (-1, -1)]

    for i, j in directions:
      neighbor = (point[0] + i, point[1] + j)
      if 0 <= neighbor[0] < img.shape[0] and 0 <= neighbor[1] < img.shape[1]:
        yield neighbor

  def visible_tiles_from_tile(self, tile):
    x, y = tile
    tiles = []

    theta = 0
    for degrees in np.arange(0, 360, 5):
      theta = degrees * np.pi / 180.0

      x1 = int(radius / self.tile_size * np.cos(theta) + x)
      y1 = int(radius / self.tile_size * np.sin(theta) + y)

      line = get_bresenham_line(x, y, x1, y1, self.tile_is_filled)
      tiles.extend(line)

    return list(set(tiles))

  def as_pixel_graph(self):
    G = nx.Graph()

    print 'generating pixel graph...   '

    for x in range(self.img.shape[0]):
      for y in range(self.img.shape[1]):
        filled = not self.img[x][y].all()

        if filled:
          continue

        G.add_node((x, y))
        G.add_edges_from([
          ((x, y), neighbor)
          for neighbor in self.neighbors_for_point((x, y))
          if self.img[neighbor[0]][neighbor[1]].all()  # no filled neighbors
        ])

    print 'done.'

    return G

def precompute_visible_points(map):
  visible = {}
  for i in xrange(map.tile_height):
    print 'row %s / %s' % (i, map.tile_height)

    for j in xrange(map.tile_width):
      visible[(i, j)] = map.visible_tiles_from_tile((i, j))
      print '\tcolumn %s / %s' % (j, map.tile_width)

  return visible

map = Map(img)
start_tile = np.array(start) / map.tile_size

visible = None
try:
  print 'loading visible points...'
  with open('map-visible-points.p') as f:
    visible = pickle.load(f)

  print 'done.'
except Exception as e:
  print e
  visible = precompute_visible_points(map)
  with open('map-visible-points.p', 'w') as f:
    pickle.dump(visible, f)

coverage = 0.999
tiles = filter(lambda tile: not map.tile_is_filled(tile), map.tiles())
min_coverage = coverage * len(tiles)

C = set()
paths = []

while len(C) < min_coverage:
  print 'len(C): %s / %s' % (len(C), min_coverage)

  subsets = []

  for tile in set(tiles) - C:

    S = set(visible[tile])
    if paths:
      path = [tile]
    else:
      path = [tile]

    cost = float(len(path))
    addition = len(S - C)
    alpha = cost / addition if addition else 1000

    heapq.heappush(subsets, (alpha, S, path))

  _, S, path = heapq.heappop(subsets)

  print 'addition: %s from: %s' % (len(S - C), path[0])

  C = C | S
  paths.append(path)

  print '\tnew total: %s' % (len(C) / float(len(tiles)))

patrol = np.array([
  map.point_for_tile(tile)
  for tile in np.concatenate(paths)
  if not map.point_is_filled(map.point_for_tile(tile))
])

G = map.as_pixel_graph()


tpatrol = [tuple(e) for e in patrol]
pairs = list(combinations(tpatrol, 2))

shortest = {}
print 'precomputing shortest paths...   (takes about a minute)'
for start, end in pairs:
  path = nx.shortest_path(G, start, end)
  shortest[(start, end)] = shortest[(end, start)] = path
print 'done.'

matrix = [
  [
    len(shortest[(tpatrol[i], tpatrol[j])])
    if i != j
    else 0
    for j in range(len(tpatrol))
  ]

  for i in range(len(tpatrol))
]

matrix_sym = pytsp.atsp_tsp(matrix, strategy="avg")
outf = "/tmp/myroute.tsp"
with open(outf, 'w') as dest:
  dest.write(pytsp.dumps_matrix(matrix_sym, name="My Route"))

tour = pytsp.run(outf, solver="LKH")
waypoints = [tpatrol[i] for i in tour['tour']]
edges = list(izip(waypoints[:-1], waypoints[1:]))

patrol = np.concatenate([
  shortest[edge] for edge in edges
])

######################
# Display the path
# NOTE: COLOR FILL DOES FOLLOW LINE OF SIGHT, just a quick helpful visualization
plt.imshow(img)
plt.plot(patrol[:,0], patrol[:,1], 'g--')
plt.plot(patrol[:,0], patrol[:,1], 'g', linewidth=2*radius, solid_capstyle='round', alpha=.2)

plt.xlim([0,img.shape[1]])
plt.ylim([img.shape[0],0])
plt.show()