"""Pick cells with a battery: at most `battery` moves ('C' costs nothing). Heuristic."""
from collections import deque

import pick_cells
from dirt_common import best_subset, neighbours, parse, route_with_cleaning

NEAR = 6        # look this many steps around the robot first; search further only if needed


def moves_used(route):
    return sum(ch != 'C' for ch in route)


def solution(room, capacity, battery):
    route = pick_cells.solution(room, capacity)
    if moves_used(route) <= battery:
        return route                            # the unlimited best already fits
    start, dirt = parse(room)
    best_route, best_value = '', 0
    for alpha in (1, 1.5, 2):
        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):
    """Walk towards the most useful dirt per step, then clean the best subset of what was visited."""
    limit = min(capacity, sum(dirt.values()))
    mask = (1 << (limit + 1)) - 1
    reach = 1                                   # bit s set: some visited cells add up to exactly s
    pos, left = start, battery
    visited = {start}
    moves, order = [], [(start, 0)]

    while left > 0:
        have = reach.bit_length() - 1
        if have == limit:
            break
        # how much one more cell with dirt w would raise the best amount we could collect
        gain_of = [0] + [((reach | (reach << w)) & mask).bit_length() - 1 - have for w in range(1, 6)]
        target, parent = find_target(pos, dirt, visited, gain_of, limit - have, min(left, NEAR), alpha)
        if target is None and left > NEAR:
            target, parent = find_target(pos, dirt, visited, gain_of, limit - have, left, alpha)
        if target is None:
            break
        path, cell = [], target
        while cell != pos:
            prev, go = parent[cell]
            path.append((go, cell))
            cell = prev
        for go, cell in reversed(path):
            moves.append(go)
            left -= 1
            if cell not in visited:
                visited.add(cell)
                order.append((cell, len(moves)))
                reach |= (reach << dirt[cell]) & mask
        pos = target

    dirty = [cell for cell, _ in order if dirt[cell] > 0]
    value, picked = best_subset([dirt[cell] for cell in dirty], capacity)
    return route_with_cleaning(moves, order, {dirty[i] for i in picked}), value


def find_target(pos, dirt, visited, gain_of, still_fits, radius, alpha):
    """BFS up to `radius` steps. Among shortest paths, keep the one passing the most useful
    unvisited dirt, and return the cell with the best useful-dirt-per-step score."""
    dist, gain, parent = {pos: 0}, {pos: 0}, {}
    queue = deque([pos])
    while queue:
        cell = queue.popleft()
        if dist[cell] == radius:
            continue
        for n, go, _ in neighbours(cell, dirt):
            g = gain[cell] + (0 if n in visited else gain_of[dirt[n]])
            if n not in dist:
                dist[n], gain[n], parent[n] = dist[cell] + 1, g, (cell, go)
                queue.append(n)
            elif dist[n] == dist[cell] + 1 and g > gain[n]:
                gain[n], parent[n] = g, (cell, go)

    target, best_score = None, 0
    for cell, g in gain.items():
        if g > 0:
            score = min(g, still_fits) / dist[cell] ** alpha
            if score > best_score:
                target, best_score = cell, score
    return target, parent