"""Part 10: emptying the bag, by column generation.

Part 3's bag problem, pick-cells rule: every dirty field must be cleaned, the bag holds `capacity`,
the dock empties it, and the robot may cross dirt with its brushes up, so a trip is just a list of
fields in the order it cleans them. A plan is a set of trips from the dock: closed ones that come
back, and one open trip last that needn't. Its cost is the moves: the walk from the dock through
the trip's fields (and back, for a closed trip).

As a linear program over trips, the set-covering model:

    minimise    sum over trips r of cost_r * x_r
    such that   sum of x_r over the trips that clean field i  >= 1     for every dirty field i
                sum of x_r over the open trips                  = 1
                x_r >= 0

(Covering instead of partitioning loses nothing: a field cleaned twice can be dropped from all
but one trip, and with shortest-path walking distances that never makes a trip longer.)

There are more trips than atoms in the universe, so the program starts with a few, one per field,
and **column generation** adds the rest only when they help. After solving the program over the
trips so far, each field i has a dual price pi_i and the open-trip row a price sigma. A trip is
worth adding only if its **reduced cost**, cost - sum of pi over its fields (- sigma if open), is
negative. Finding the most negative one is the **pricing problem**: a shortest path from the dock
through fields whose "length" is walking distance minus price, with the bag's capacity as a
resource and no field twice. `price` solves it with labels and dominance. When no trip has a
negative reduced cost, the program's value is its true minimum over *all* trips: a lower bound on
every plan. Then an integer program over the trips generated (HiGHS) gives a real plan, whose
distance from the bound says how good it is: **price and branch**.
"""
import math
from collections import deque

from dirt_common import neighbours, parse

EPS = 1e-9
NG = 8                  # how many nearest fields a label remembers visiting (ng-routes)


class Instance:
    """The dirty fields, their dirt and the walking distances between them and the dock."""

    def __init__(self, room, capacity):
        self.dock, dirt = parse(room)
        self.room, self.capacity = room, capacity
        self.fields = sorted(f for f in dirt if dirt[f] > 0)
        self.dirt = [dirt[f] for f in self.fields]
        self.n = len(self.fields)
        points = self.fields + [self.dock]              # the dock is point n
        self.dist = []
        for p in points:
            d, queue = {p: 0}, deque([p])
            while queue:
                cell = queue.popleft()
                for nb, _, _ in neighbours(cell, dirt):
                    if nb not in d:
                        d[nb] = d[cell] + 1
                        queue.append(nb)
            self.dist.append([d[q] for q in points])
        # Each field's ng-neighbourhood: itself and its NG nearest dirty fields (see price).
        self.ng = []
        for j in range(self.n):
            near = sorted(range(self.n), key=lambda k: (self.dist[j][k], k))[:NG + 1]
            self.ng.append(sum(1 << k for k in near) | 1 << j)

    def cost(self, trip, open_):
        """Moves for a trip: from the dock through the fields in order, and home if closed."""
        n, d = self.n, self.dist
        walk = d[n][trip[0]] + sum(d[a][b] for a, b in zip(trip, trip[1:]))
        return walk if open_ else walk + d[trip[-1]][n]


def price(inst, pi, sigma, limit=30, max_labels=5_000_000):
    """The trips with the most negative reduced cost, by labels and dominance.

    A label is a trip under construction: (last field, load, memory, reduced cost so far, fields
    in order). Extending it to a field j that fits costs dist(last, j) - pi_j. A label is dropped
    when another at the same field has no more load, no more reduced cost and remembers no field
    it doesn't: anything the dropped one can still become, the other can become at least as
    cheaply.

    Remembering every field visited makes labels almost never dominate each other, and the search
    explodes. So a label remembers only the fields near where it is, its ng-memory (Baldacci,
    Mingozzi and Roberti): extending to j keeps the memory's fields that are in j's
    neighbourhood, plus j. A trip may then come back to a field it visited long ago, so the
    labels describe more trips than the elementary ones (ng-routes), and when none of them has a
    negative reduced cost, no elementary trip has either: the bound stays valid. A trip that comes
    back to a field is still a real plan, only a wasteful one (the robot drives over a field it
    has already cleaned), so it goes into the program as it is, covering that field as often as
    it visits it: exactly the count its reduced cost was computed with.

    Returns up to `limit` trips as (reduced cost, fields in order, open), most negative first, or
    None if the labels ran past max_labels."""
    n, d, Q, ng = inst.n, inst.dist, inst.capacity, inst.ng
    found = []
    buckets = [[] for _ in range(n)]                    # the undominated labels at each field
    queue = deque()
    for j in range(n):
        label = (j, inst.dirt[j], 1 << j, d[n][j] - pi[j], (j,))
        buckets[j].append(label)
        queue.append(label)
    made = n
    while queue:
        label = queue.popleft()
        last, load, memory, red, trip = label
        if label not in buckets[last]:
            continue                                    # dominated after it was queued
        closed = red + d[last][n]
        opened = red - sigma
        if closed < -EPS:
            found.append((closed, trip, False))
        if opened < -EPS:
            found.append((opened, trip, True))
        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, red + d[last][j] - pi[j], trip + (j,))
            keep, bucket = True, buckets[j]
            for other in bucket:
                if other[1] <= new[1] and other[3] <= new[3] + EPS and other[2] & new[2] == other[2]:
                    keep = False
                    break
            if not keep:
                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:
                return None
    found.sort(key=lambda t: t[0])
    unique, keys = [], set()
    for red, trip, open_ in found:
        key = (tuple(sorted(trip)), open_)
        if key not in keys:                             # one order per set of visits is enough
            keys.add(key)
            unique.append((red, trip, open_))
        if len(unique) == limit:
            break
    return unique


class Master:
    """The covering program over the trips generated so far, in HiGHS."""

    def __init__(self, inst):
        import highspy
        import numpy as np
        self.np, self.inst = np, inst
        self.h = h = highspy.Highs()
        h.setOptionValue('output_flag', False)
        n = inst.n
        for i in range(n):                              # rows 0..n-1: every field covered
            h.addRow(1.0, np.inf, 0, np.array([], dtype=np.int32), np.array([]))
        h.addRow(1.0, 1.0, 0, np.array([], dtype=np.int32), np.array([]))    # row n: one open trip
        self.trips = []                                 # (fields in order, open, cost)

    def add(self, trip, open_):
        np, n = self.np, self.inst.n
        cost = self.inst.cost(trip, open_)
        visits = {}
        for j in trip:                                  # a field visited twice is covered twice
            visits[j] = visits.get(j, 0) + 1
        if open_:
            visits[n] = 1
        rows = sorted(visits)
        self.h.addCol(float(cost), 0.0, np.inf, len(rows), np.array(rows, dtype=np.int32),
                      np.array([float(visits[r]) for r in rows]))
        self.trips.append((tuple(trip), open_, cost))

    def solve(self):
        self.h.run()
        sol = self.h.getSolution()
        dual = list(sol.row_dual)
        return self.h.getInfo().objective_function_value, dual[:self.inst.n], dual[self.inst.n], list(sol.col_value)


class ExactMaster:
    """The same covering program in exact fractions, solved from scratch by Part 8's simplex.py
    every round. Only for small rooms; the dual prices it gives are exact."""

    def __init__(self, inst):
        self.inst, self.trips = inst, []

    def add(self, trip, open_):
        self.trips.append((tuple(trip), open_, self.inst.cost(trip, open_)))

    def solve(self):
        import simplex
        n = self.inst.n
        rows = [({c: trip.count(i) for c, (trip, _, _) in enumerate(self.trips) if i in trip}, simplex.GE, 1) for i in range(n)]
        rows.append(({c: 1 for c, (_, open_, _) in enumerate(self.trips) if open_}, simplex.EQ, 1))
        result = simplex.solve([cost for _, _, cost in self.trips], rows)
        return result.value, result.y[:n], result.y[n], result.x


def generate(room, capacity, record=None, max_rounds=10_000, solver='highs'):
    """Column generation to the end, with the master in HiGHS or ('exact') in simplex.py. Returns
    (bound, master, rounds, instance); record(round, value, prices, sigma, added) is called after
    every round."""
    inst = Instance(room, capacity)
    master = Master(inst) if solver == 'highs' else ExactMaster(inst)
    for j in range(inst.n):
        master.add((j,), False)
        master.add((j,), True)
    for rounds in range(1, max_rounds + 1):
        value, pi, sigma, _ = master.solve()
        new = price(inst, pi, sigma)
        if record:
            record(rounds, value, pi, sigma, new)
        if not new:
            return value, master, rounds, inst
        for _, trip, open_ in new:
            master.add(trip, open_)
    raise RuntimeError('column generation did not settle')


def completion(inst, pi, sigma):
    """A lower bound on how much a trip's reduced cost can still fall from each field with r units
    of room left in the bag: the cheapest way to finish, allowing fields to repeat (a q-route
    relaxation), by dynamic programming over (field, room left). Returns g[r][j]."""
    n, d, Q = inst.n, inst.dist, inst.capacity
    finish = [min(d[j][n], -sigma) for j in range(n)]           # go home, or stop (open trip)
    g = [finish[:] for _ in range(Q + 1)]
    for r in range(1, Q + 1):
        for j in range(n):
            best = finish[j]
            for k in range(n):
                if k != j and inst.dirt[k] <= r:
                    best = min(best, d[j][k] - pi[k] + g[r - inst.dirt[k]][k])
            g[r][j] = best
    return g


def enumerate_trips(inst, pi, sigma, gap, max_trips=500_000):
    """Every trip whose reduced cost is at most gap, one cheapest order per set of fields.

    No plan costing at most (bound + gap) can use a trip with a larger reduced cost, because a
    plan's cost is the LP bound plus the reduced costs of its trips plus a sum that is never
    negative. Partial trips are cut off with `completion`, which never overestimates."""
    n, d, Q = inst.n, inst.dist, inst.capacity
    g = completion(inst, pi, sigma)
    best = {}                                           # (fields, open) -> (cost, order)
    stack = [((j,), inst.dirt[j], 1 << j, d[n][j] - pi[j]) for j in range(n)]
    while stack:
        trip, load, seen, red = stack.pop()
        last = trip[-1]
        if red + g[Q - load][last] > gap + EPS:
            continue
        for red_end, open_ in ((red + d[last][n], False), (red - sigma, True)):
            if red_end <= gap + EPS:
                cost = inst.cost(trip, open_)
                key = (seen, open_)
                if key not in best or cost < best[key][0]:
                    best[key] = (cost, trip)
                    if len(best) > max_trips:
                        raise RuntimeError('too many trips to enumerate')
        for j in range(n):
            if not seen >> j & 1 and load + inst.dirt[j] <= Q:
                stack.append((trip + (j,), load + inst.dirt[j], seen | 1 << j, red + d[last][j] - pi[j]))
    return [(trip, open_) for (seen, open_), (cost, trip) in best.items()]


def solve(room, capacity, time_limit=600):
    """The best plan, proven. Column generation gives the bound; Part 3's fast plan and the
    integer program over the generated trips give a first plan; if they don't meet, every trip
    that could still be in a better plan is enumerated and the integer program solved over them.
    Returns a dict: commands, moves, bound (the LP's), proven, trips enumerated, rounds."""
    import highspy
    import pick_cells_dock
    from test_dock import simulate_dock
    bound, master, rounds, inst = generate(room, capacity)
    lower = math.ceil(bound - 1e-6)
    value, pi, sigma, _ = master.solve()
    fast = pick_cells_dock.solution(room, capacity)
    best_cmds, best, plan = fast, simulate_dock(room, fast, capacity, True), None
    first = integer(master, time_limit)
    if first and first[1] < best:
        best_cmds, best, plan = first
    enumerated = 0
    proven = best == lower
    if not proven:
        try:
            trips = enumerate_trips(inst, pi, sigma, best - bound)
        except RuntimeError:                            # the gap is too wide to enumerate
            trips = None
        if trips is not None:
            enumerated = len(trips)
            full = Master(inst)
            for trip, open_ in trips:
                full.add(trip, open_)
            solved = integer(full, time_limit)
            if solved and solved[1] < best:
                best_cmds, best, plan = solved
            proven = solved is not None
    return {'commands': best_cmds, 'moves': best, 'bound': bound, 'proven': proven, 'plan': plan,
            'enumerated': enumerated, 'rounds': rounds, 'columns': len(master.trips)}


def integer(master, time_limit):
    """The integer program over a master's trips. Returns (commands, moves, trips) or None."""
    import highspy
    h, np = master.h, master.np
    k = len(master.trips)
    if not k:
        return None
    h.setOptionValue('time_limit', float(time_limit))
    h.changeColsIntegrality(k, np.arange(k, dtype=np.int32), np.array([highspy.HighsVarType.kInteger] * k))
    h.run()
    if h.getModelStatus() != highspy.HighsModelStatus.kOptimal:
        return None
    chosen = [master.trips[c] for c, v in enumerate(h.getSolution().col_value) if v > 0.5]
    cmds = commands(master.inst, chosen)
    return cmds, sum(1 for ch in cmds if ch != 'C'), chosen   # skipping a repeated visit only helps


def commands(inst, trips):
    """Turn trips into the robot's commands: closed trips first, the open one last. A field in
    two trips is cleaned by the first."""
    dock, dirt_cells = inst.dock, set(parse(inst.room)[1])
    order = sorted(trips, key=lambda t: t[1])           # closed (False) before open (True)
    done, out, at = set(), [], dock
    for trip, open_, _ in order:
        stops = []
        for j in trip:                                  # each field only at its first visit
            if j not in done:
                done.add(j)
                stops.append(inst.fields[j])
        for f in stops + ([] if open_ else [dock]):
            out += path(at, f, dirt_cells)
            at = f
            if f != dock:
                out.append('C')
    return ''.join(out)


def path(a, b, cells):
    """Commands along a shortest walk from a to b (brushes up: dirt doesn't matter)."""
    go = {(-1, 0): '^', (1, 0): 'v', (0, -1): '<', (0, 1): '>'}
    parent, queue = {a: None}, deque([a])
    while b not in parent:
        u = queue.popleft()
        for dr, dc in go:
            v = (u[0] + dr, u[1] + dc)
            if v in cells and v not in parent:
                parent[v] = u
                queue.append(v)
    steps = []
    while b != a:
        p = parent[b]
        steps.append(go[(b[0] - p[0], b[1] - p[1])])
        b = p
    return steps[::-1]