"""Helpers for emptying the bag: the robot's start is a dock where the bag empties.

Both pickup rules share one idea: put the dirty fields in a good order, then cut that order into
trips from the dock in the best possible way.
"""
from collections import deque

from dirt_common import neighbours

STEP_OF = {(-1, 0): '^', (1, 0): 'v', (0, -1): '<', (0, 1): '>'}


def direction(a, b):
    return STEP_OF[(b[0] - a[0], b[1] - a[1])]


def bfs_tree(dock, dirt, rng=None):
    """Distance of every field from the dock and a breadth-first tree (each field's parent).

    With rng the neighbours are explored in random order, which gives a different tree."""
    dist, parent = {dock: 0}, {dock: None}
    queue = deque([dock])
    while queue:
        cell = queue.popleft()
        nbrs = [n for n, _, _ in neighbours(cell, dirt)]
        if rng:
            rng.shuffle(nbrs)
        for n in nbrs:
            if n not in dist:
                dist[n], parent[n] = dist[cell] + 1, cell
                queue.append(n)
    return dist, parent


def tree_order(dock, dirt, dist, parent, rng=None):
    """Dirty fields in preorder of the tree, visiting the child that leads deepest last.

    Every field comes after its parent, so under 'take all' any cut of this order is legal."""
    children = {cell: [] for cell in parent}
    for cell, p in parent.items():
        if p is not None:
            children[p].append(cell)
    deepest = dict(dist)
    for cell in sorted(parent, key=dist.get, reverse=True):
        if parent[cell] is not None:
            deepest[parent[cell]] = max(deepest[parent[cell]], deepest[cell])
    order, stack = [], [dock]
    while stack:
        cell = stack.pop()
        if dirt[cell] > 0:
            order.append(cell)
        stack.extend(sorted(children[cell], key=lambda k: (-deepest[k], rng.random() if rng else 0)))
    return order


def split(weights, d0, gaps, capacity):
    """Cut fields 0..n-1 (in this order) into consecutive trips with the fewest moves.

    A closed trip over fields a..b costs d0[a] + gaps a..b + d0[b]; the last trip doesn't come
    back. gaps[k] is the walk from field k to field k+1. Returns (moves, trips as index lists).
    Uses a sliding-window minimum, so it runs in O(n)."""
    n = len(weights)
    if n == 0:
        return 0, []
    load, walked = [0] * (n + 1), [0] * n       # prefix sums of weights and of gaps
    for k in range(n):
        load[k + 1] = load[k] + weights[k]
        if k:
            walked[k] = walked[k - 1] + gaps[k - 1]
    best = [0] * (n + 1)                        # best[t]: fewest moves to finish fields 0..t-1
    came_from = [0] * (n + 1)
    key = [0] * n                               # cost of a trip starting at a, minus its gaps so far
    window = deque()
    for b in range(n):
        key[b] = best[b] + d0[b] - walked[b]
        while window and key[window[-1]] >= key[b]:
            window.pop()
        window.append(b)
        while load[b + 1] - load[window[0]] > capacity:
            window.popleft()
        a = window[0]
        best[b + 1] = key[a] + walked[b] + d0[b]
        came_from[b + 1] = a
    a = window[0]                               # the open last trip ends at field n-1
    trips, t = [list(range(a, n))], a
    while t > 0:
        s = came_from[t]
        trips.append(list(range(s, t)))
        t = s
    return key[a] + walked[n - 1], trips[::-1]


def split_reference(weights, d0, gaps, capacity):
    """The same cut by a plain double loop, for testing."""
    n = len(weights)
    if n == 0:
        return 0
    walked = [0] * n
    for k in range(1, n):
        walked[k] = walked[k - 1] + gaps[k - 1]
    inf = float('inf')
    best = [0] + [inf] * n
    for b in range(n):
        for a in range(b, -1, -1):
            if sum(weights[a:b + 1]) > capacity:
                break
            best[b + 1] = min(best[b + 1], best[a] + d0[a] + walked[b] - walked[a] + d0[b])
    return min(best[a] + d0[a] + walked[n - 1] - walked[a]
               for a in range(n) if sum(weights[a:]) <= capacity)


def tree_path(parent, cell):
    """Fields from cell up to the dock along the tree: [cell, parent, ..., dock]."""
    path = [cell]
    while parent[path[-1]] is not None:
        path.append(parent[path[-1]])
    return path


def build_route(dock, trips, parent, gap_path, pick):
    """Commands for the trips: out along the tree to the first field, gap paths between fields,
    back along the tree (except after the last trip). `gap_path(a, b, cleaned)` gives the fields
    walked from a to b, ending with b."""
    out, cleaned, pos = [], {dock}, dock
    for t, trip in enumerate(trips):
        for k, field in enumerate(trip):
            walk = reversed(tree_path(parent, field)[:-1]) if k == 0 else gap_path(pos, field, cleaned)
            for cell in walk:
                out.append(direction(pos, cell))
                pos = cell
                if not pick:
                    cleaned.add(cell)
            if pick:
                out.append('C')
            cleaned.add(field)
        if t < len(trips) - 1:
            for cell in tree_path(parent, pos)[1:]:
                out.append(direction(pos, cell))
                pos = cell
    return ''.join(out)


def bounds(dirt, dist, capacity):
    """(lower, upper) on the moves any plan needs.

    upper: one field per trip, farthest trip last.
    lower: a closed trip costs at least twice its farthest field, and every trip but the last
    ends by stepping onto the dock."""
    fields = [f for f in dirt if dirt[f] > 0]
    if not fields:
        return 0, 0
    far = max(dist[f] for f in fields)
    total = sum(dirt[f] for f in fields)
    upper = 2 * sum(dist[f] for f in fields) - far
    lower = max(-(-2 * sum(dirt[f] * dist[f] for f in fields) // capacity) - far,
                len(fields) + -(-total // capacity) - 1)
    return lower, upper