"""Parts 4 and 6: the fewest commands on any surface, exactly, for small rooms.

Part 3's exact solver searched over (robot's field, set of cleaned fields), breadth first, and cut
the search short with its lower bound, the checkerboard part included. That part is only true
when the floor can be coloured like a chessboard. This version takes the colouring as an
argument, so the same search runs Part 3's way (colour by (row + column) mod 2, trusted), without
a colouring, or with a proper one from two_colour.py.
"""
from surface import DIRS, cross, depth_first, robot_dir, simulate


def floor_from(surface, start):
    seen, todo = [start], [start]
    known = {start}
    while todo:
        f = todo.pop()
        for g in surface.glue[f]:
            if g is not None and g[0] not in known:
                known.add(g[0])
                seen.append(g[0])
                todo.append(g[0])
    return seen


def lower_bound(surface, start, colour=None):
    """Part 3's three facts. `colour` maps field -> 0 or 1; with None only 'one new field per
    command' and 'dead ends' are used."""
    floor = floor_from(surface, start)
    n = len(floor)
    dead_ends = sum(1 for f in floor
                    if f != start and sum(g is not None for g in surface.glue[f]) == 1)
    bound = max(n - 1, n - 1 + max(0, dead_ends - 1))
    if colour is not None:
        same = sum(colour[f] == colour[start] for f in floor)
        bound = max(bound, 2 * same - 2, 2 * (n - same) - 1)
    return bound


def trimmed_walk(surface, start, frame=(0, False)):
    """Part 1's walk, cut off after the last new field."""
    cmds = depth_first(surface, start, frame)
    field, seen, last = start, {start}, 0
    for i, ch in enumerate(cmds):
        _, field, frame = simulate(surface, field, ch, frame)
        if field not in seen:
            seen.add(field)
            last = i + 1
    return cmds[:last]


def fewest(surface, start, colour=None, prune=True, frame=(0, False), max_states=3_000_000):
    """The fewest commands that clean every field reachable from start.

    Breadth-first search over (field, cleaned set), one layer per command. With prune, a state
    is dropped when it can't beat the best route known so far: it still needs one step per
    field left, and with a colouring Part 3's alternating-colours count too. Without prune it
    is a plain brute force, the referee. Returns the commands, or None past max_states.
    """
    floor = floor_from(surface, start)
    index = {f: i for i, f in enumerate(floor)}
    full = (1 << len(floor)) - 1
    moves = [[(index[g[0]], side) for side, g in enumerate(surface.glue[f]) if g is not None]
             for f in floor]
    black = 0
    if colour is not None:
        black = sum(1 << i for i, f in enumerate(floor) if colour[f])
    best = trimmed_walk(surface, start, frame) if prune else None
    limit = len(best) if best is not None else 2 * len(floor)

    first = (0, 1)
    if first[1] == full:
        return ''
    parent = {first: None}
    layer = [first]
    for k in range(1, limit):
        nxt = []
        for state in layer:
            i, mask = state
            for j, side in moves[i]:
                new = (j, mask | 1 << j)
                if new in parent:
                    continue
                if prune:
                    left = full & ~new[1]
                    need = bin(left).count('1')
                    if colour is not None:
                        same = bin(left & (black if black >> j & 1 else ~black)).count('1')
                        other = need - same
                        need = max(need, 2 * same, 2 * other - 1)
                    if k + need >= limit:
                        continue
                parent[new] = (state, side)
                if new[1] == full:
                    return commands(surface, floor, parent, new, frame)
                nxt.append(new)
        if len(parent) > max_states:
            return None
        layer = nxt
    return best                             # nothing shorter than the walk exists


def commands(surface, floor, parent, state, frame):
    """Read the route back as sides, then replay it to translate each side into the robot's own
    direction for the frame it has at that moment."""
    sides = []
    while parent[state] is not None:
        state, side = parent[state]
        sides.append(side)
    field, out = floor[0], []
    for side in reversed(sides):
        out.append(DIRS[robot_dir(frame, side)])
        field, frame = cross(surface, field, frame, side)
    return ''.join(out)