"""Part 11: the battery, bounded by a price on every move (Lagrangian relaxation).

Part 2's pick-cells robot with a battery: at most `battery` moves (cleaning is free), a bag of
`capacity`, and as much dirt as possible. As in Part 10, a route is a trip: the fields it cleans,
in order, walking the shortest way between them, and it may end anywhere. Its load is the dirt it
cleans and its moves are the walk from the dock (column_gen.Instance.cost, open).

Two budgets make it hard: the bag (a knapsack) and the battery (orienteering). Price every move at
lam units of dirt and drop the battery:

    L(lam) = lam * battery + max over trips with load <= capacity of (load - lam * moves)

For any lam >= 0 and any trip within the battery, load <= load + lam * (battery - moves) <= L(lam),
so every price gives an upper bound: **weak duality**. The best price is the one with the lowest
L: the **Lagrangian dual**. L is the maximum of one straight line per trip, so it is convex and
piecewise linear, and its slope at lam is battery - moves of the trip that is best there: a
**subgradient**, which says which way to move the price.

The trip that is best at a price (the **inner problem**) is Part 10's pricing problem with every
field's dirt as its prize and every move costing lam: a path from the dock under one resource, the
bag, solved with the same labels and dominance. With `memory` set to a number, a label remembers
only that many nearest fields (ng-routes, as in Part 10): the trips it finds may come back to a
field far behind and clean it again, which can only raise L, so every bound stays valid.
"""
from collections import deque
from fractions import Fraction

import column_gen

EPS = 1e-9


class Instance(column_gen.Instance):
    """Part 10's instance (the dirty fields, their dirt, walking distances, the dock as point n),
    with each field's memory: its `memory` nearest fields, or every field (elementary trips)."""

    def __init__(self, room, capacity, memory=None):
        super().__init__(room, capacity)
        n = self.n
        if memory is None or memory >= n:
            self.ng = [(1 << n) - 1] * n
        else:
            self.ng = []
            for j in range(n):
                near = sorted(range(n), key=lambda k: (self.dist[j][k], k))[:memory + 1]
                self.ng.append(sum(1 << k for k in near) | 1 << j)

    def moves(self, trip):
        return self.cost(trip, True) if trip else 0

    def load(self, trip):
        """The dirt a trip really cleans: a field it comes back to is clean already."""
        return sum(self.dirt[j] for j in set(trip))


def best_trip(inst, lam, max_labels=5_000_000):
    """The trip with the largest load - lam * moves, the bag respected and no battery.

    A label is a trip under construction: (last field, load, memory, value so far, moves, fields in
    order). Extending it to field j adds dirt_j - lam * dist(last, j). A label is dropped when
    another at the same field has no more load, no less value and remembers no field it doesn't:
    whatever the dropped one can still become, the other can become at least as well. It is also
    dropped when even filling the bag for free couldn't beat the best trip found: its value plus
    the room left in the bag. The best starts as the better of the empty trip (value 0) and Part
    2's answer without a battery, which matters at low prices, where every load ties with every
    way of walking it and the labels hardly dominate each other. Returns (value, fields in order,
    moves, load)."""
    n, d, Q, ng = inst.n, inst.dist, inst.capacity, inst.ng
    best = (0, (), 0, 0)
    full = unlimited(inst)
    if full and inst.load(full) - lam * inst.moves(full) > EPS:
        best = (inst.load(full) - lam * inst.moves(full), full, inst.moves(full), inst.load(full))
    buckets = [[] for _ in range(n)]
    queue = deque()
    for j in range(n):
        if inst.dirt[j] <= Q:
            label = (j, inst.dirt[j], 1 << j, inst.dirt[j] - lam * d[n][j], d[n][j], (j,))
            buckets[j].append(label)
            queue.append(label)
    made = len(queue)
    while queue:
        label = queue.popleft()
        last, load, memory, value, moves, trip = label
        if label not in buckets[last]:
            continue                                    # dominated after it was queued
        if value > best[0] + EPS:
            best = (value, trip, moves, load)
        for j in range(n):
            if memory >> j & 1 or load + inst.dirt[j] > Q:
                continue
            new = (j, load + inst.dirt[j], (memory & ng[j]) | 1 << j, value + inst.dirt[j] - lam * d[last][j],
                   moves + d[last][j], trip + (j,))
            if new[3] + Q - new[1] <= best[0] + EPS:
                continue                                # can't beat the best even with a free full bag
            bucket = buckets[j]
            if any(o[1] <= new[1] and o[3] >= new[3] - EPS and o[2] & new[2] == o[2] for o in bucket):
                continue
            bucket[:] = [o for o in bucket if not (new[1] <= o[1] and new[3] >= o[3] - EPS and new[2] & o[2] == new[2])]
            bucket.append(new)
            queue.append(new)
            made += 1
            if made > max_labels:
                raise RuntimeError('too many labels')
    return best


def within(inst, trip, battery):
    """The trip cut where the battery runs out, each field cleaned once: a real route."""
    n, d = inst.n, inst.dist
    at, moves, kept = n, 0, []
    for j in trip:
        moves += d[at][j]
        if moves > battery:
            break
        if j not in kept:
            kept.append(j)
        at = j
    return tuple(kept)


def subgradient(inst, battery, start, steps=60, record=None):
    """Lower the bound by subgradient steps on the price (Held, Wolfe and Crowder; Polyak's step).

    `start` is a route within the battery, as a trip; it and every inner trip cut to the battery
    are the best route known, `low`. The price starts at what a move earned on `start`. Each step:
    solve the inner problem, L = lam * battery + value; the slope is g = battery - moves; then

        lam  <-  max(0, lam - theta * (L - low) / g^2 * g)

    which would land exactly on `low` if L were the straight line through this step. theta starts at
    1 and halves after three steps that don't lower the best bound. It stops when the bound proves
    `low` (loads are whole numbers, so best bound < low + 1 is enough), when the slope is 0 (then
    this price is the best there is), or after `steps` steps.

    Returns (best bound, best route as a trip, steps taken); record(step, lam, L, trip, moves, load,
    low) is called after every step."""
    low_trip = start
    low = inst.load(start)
    lam = low / battery if battery else 1.0
    theta, best, stall = 1.0, float('inf'), 0
    for k in range(1, steps + 1):
        value, trip, moves, load = best_trip(inst, lam)
        bound = lam * battery + value
        cut = within(inst, trip, battery)
        if inst.load(cut) > low:
            low_trip, low = cut, inst.load(cut)
        if bound < best - EPS:
            best, stall = bound, 0
        else:
            stall += 1
            if stall == 3:
                theta, stall = theta / 2, 0
        if record:
            record(k, lam, bound, trip, moves, load, low)
        g = battery - moves
        if best < low + 1 - 1e-6 or g == 0:
            return best, low_trip, k
        lam = max(0.0, lam - theta * (bound - low) / (g * g) * g)
    return best, low_trip, steps


def trip_of(inst, commands):
    """The fields a pick-cells route cleans, in order: the route as a trip. The trip walks the
    shortest way between them, so it never needs more moves than the route."""
    step = {'^': (-1, 0), 'v': (1, 0), '<': (0, -1), '>': (0, 1)}
    index = {f: j for j, f in enumerate(inst.fields)}
    at, trip = inst.dock, []
    for ch in commands:
        if ch == 'C':
            trip.append(index[at])
        else:
            at = (at[0] + step[ch][0], at[1] + step[ch][1])
    return tuple(trip)


def unlimited(inst):
    """Part 2's answer without a battery, as a trip: its walk visits every field, and subset sum
    picks what to clean. No trip has more dirt, so at price 0 it is a best trip."""
    import pick_cells
    return trip_of(inst, pick_cells.solution(inst.room, inst.capacity))


def dual(inst, battery, exact=True, record=None):
    """The Lagrangian dual exactly, by cutting planes on L (Kelley's method; in one dimension each
    step also halves the room the minimum can be in).

    Every trip t gives a line lam * (battery - moves_t) + load_t below L. Keep the best line that
    rises (moves <= battery) and the best that falls (moves > battery); the dual's minimum lies
    where they cross. Solve the inner problem at the crossing: if no trip lies above both lines
    there, the crossing is the minimum; otherwise the new trip replaces one of them. It starts with
    the empty trip, which rises, and Part 2's unlimited answer, which has the most dirt there is:
    if that fits in the battery it is the answer. With exact fractions it stops after finitely many
    steps, since there are finitely many trips; exact=False works in floating point, for big rooms.

    Returns (L*, the price lam*, steps); record(step, lam, value of the lines there, trip found) is
    called after every step."""
    trip = unlimited(inst)
    moves, load = inst.moves(trip), inst.load(trip)
    if record:
        record(0, 0, load, trip)
    if moves <= battery:
        return Fraction(load), Fraction(0), 1           # with no price the best trip fits already
    rise, fall = (0, 0), (moves, load)                  # (moves, load); the empty trip: L >= lam * battery
    for k in range(1, 200):
        (m1, l1), (m2, l2) = rise, fall
        lam = Fraction(l2 - l1, m2 - m1) if exact else (l2 - l1) / (m2 - m1)    # where the two lines cross
        level = lam * (battery - m1) + l1
        value, trip, moves, load = best_trip(inst, lam)
        if record:
            record(k, lam, level, trip)
        if lam * battery + value <= level + (0 if exact else 1e-7):
            return level, lam, k + 1
        if moves <= battery:
            rise = (moves, load)
        else:
            fall = (moves, load)
    raise RuntimeError('the dual did not settle')