"""Take all with a battery: at most `battery` moves. Heuristic."""
from collections import deque

from dirt_common import neighbours, parse, walk
from take_all_fast import best_connected_set


def solution(room, capacity, battery):
    start, dirt = parse(room)
    _, cells, _ = best_connected_set(start, dirt, capacity)
    moves, order = walk(start, cells)
    full = moves[:order[-1][1]]
    if len(full) <= battery:
        return ''.join(full)                    # the unlimited answer already fits

    # candidate 1: follow that route until the battery runs out
    best_route = ''.join(full[:battery])
    best_value = sum(dirt[cell] for cell, i in order if i <= battery)
    for alpha in (1, 1.5, 2):                   # more candidates: grow greedily
        route, value = greedy(start, dirt, capacity, battery, alpha)
        if value > best_value or (value == best_value and len(route) < len(best_route)):
            best_route, best_value = route, value
    return best_route


def greedy(start, dirt, capacity, battery, alpha):
    """Repeatedly step into the uncleaned cell with the most dirt per move that still fits.

    The robot only walks over cells it has already cleaned, so reaching a cell costs
    battery but no capacity.
    """
    cleaned, pos, left, room_left = {start}, start, battery, capacity
    moves = []
    while left > 0:
        target, entry, best_score = None, None, 0
        dist, parent = {pos: 0}, {}
        queue = deque([pos])
        while queue:
            cell = queue.popleft()
            d = dist[cell]
            if d + 1 > left or 5 / (d + 1) ** alpha <= best_score:
                break                           # nothing farther away can score higher
            for n, go, _ in neighbours(cell, dirt):
                if n in cleaned:
                    if n not in dist:
                        dist[n], parent[n] = d + 1, (cell, go)
                        queue.append(n)
                elif dirt[n] <= room_left:
                    score = dirt[n] / (d + 1) ** alpha
                    if score > best_score:
                        target, entry, best_score = n, (cell, go), score
        if target is None:
            break
        parent[target] = entry
        steps, cell = [], target
        while cell != pos:
            cell, go = parent[cell]
            steps.append(go)
        steps.reverse()
        moves.extend(steps)
        left -= len(steps)
        cleaned.add(target)
        room_left -= dirt[target]
        pos = target
    return ''.join(moves), capacity - room_left