"""Emptying the bag, pick cells: the robot drives over dirt freely and cleans with 'C'.

Put the dirty fields in a good order, cut that order into trips in the best possible way, then
improve the order with small changes (move one field, reverse a stretch), re-cutting each time.
"""
import time
from array import array

from dirt_common import neighbours, parse
from dock_common import bfs_tree, build_route, split, tree_order

FAR = 65535
NEAR = 8                                        # how many nearby fields each field looks at


def distance_table(cells, dirt):
    """Walking distance between every pair of fields (one BFS per field, rows as compact arrays),
    plus each field's NEAR closest dirty fields."""
    index = {cell: i for i, cell in enumerate(cells)}
    adjacent = [[index[n] for n, _, _ in neighbours(cell, index)] for cell in cells]
    dirty = [dirt[cell] > 0 for cell in cells]
    table, near = [], []
    for s in range(len(cells)):
        dist = array('H', [FAR]) * len(cells)
        dist[s] = 0
        queue, close = [s], []
        for u in queue:                         # the list grows while we walk it: a BFS queue
            for v in adjacent[u]:
                if dist[v] == FAR:
                    dist[v] = dist[u] + 1
                    queue.append(v)
                    if dirty[v] and len(close) < NEAR:
                        close.append(cells[v])
        table.append(dist)
        near.append(close)
    return index, table, near


def solution(room, capacity, seconds=2.0):
    dock, dirt = parse(room)
    assert capacity >= 5, 'every field must fit in the bag on its own'
    dist, parent = bfs_tree(dock, dirt)
    order = tree_order(dock, dirt, dist, parent)
    if not order:
        return ''
    cells = list(dirt)
    index, table, near = distance_table(cells, dirt)
    home = table[index[dock]]

    def cut(order):
        ids = [index[f] for f in order]
        gaps = [table[ids[k]][ids[k + 1]] for k in range(len(ids) - 1)]
        return split([dirt[f] for f in order], [home[i] for i in ids], gaps, capacity)

    best, trips = cut(order)
    deadline = time.perf_counter() + seconds
    improved = True
    while improved and time.perf_counter() < deadline:
        improved = False
        for field in list(order):
            for other in near[index[field]]:
                for candidate in changes(order, field, other):
                    moves, new_trips = cut(candidate)
                    if moves < best:
                        order, best, trips, improved = candidate, moves, new_trips, True
                        break
                else:
                    continue
                break                           # the order changed: move on to the next field
            if time.perf_counter() >= deadline:
                break

    def gap_path(a, b, cleaned):
        row, path = table[index[b]], []
        while a != b:
            a = next(n for n, _, _ in neighbours(a, index) if row[index[n]] == row[index[a]] - 1)
            path.append(a)
        return path

    return build_route(dock, [[order[i] for i in trip] for trip in trips], parent, gap_path, pick=True)


def changes(order, x, y):
    """New orders that put field x right next to its nearby field y."""
    i, j = order.index(x), order.index(y)
    without = order[:i] + order[i + 1:]
    k = without.index(y)
    yield without[:k + 1] + [x] + without[k + 1:]           # move x to just after y
    yield without[:k] + [x] + without[k:]                   # move x to just before y
    if i < j:
        yield order[:i + 1] + order[i + 1:j + 1][::-1] + order[j + 1:]   # reverse so y follows x
    else:
        yield order[:j] + order[j:i][::-1] + order[i:]                   # reverse so y precedes x
