"""Part 11: walks that may leave fields out, as integer programs in HiGHS.

Part 9's program describes one route that cleans the whole floor. Two of the robot's problems need
walks that cover only part of it: Part 2's battery, where the robot can't reach everything and
chooses what to clean, and Part 3's two robots, where each cleans its share. So each walk gets one
more variable per field,

    y_v in {0, 1}    the walk reaches v (y_s = 1 on its start s)

and Part 9's rules (x_e walks per pair of neighbours, z_v where it ends, parity through w_v) hold
for what it reaches:

    x_e <= 2 y_v               a walked pair touches only reached fields
    z_v <= y_v                 it ends on a field it reaches
    x(d(v)) + z_v >= 2 y_v     a reached field is entered and left, or the walk ends there
    flow                       s sends one unit to every reached field, only along walked pairs
    sum x >= sum y - 1         one new field per move at most

The flow is Part 9's compact way to say "connected": a walk whose pieces don't hang together
can't carry it. A reached field the flow can't get to is impossible, so the walk's edges round s,
closed by a free edge from its end back to s, have every degree even and hang together: Hierholzer's
algorithm (branch_cut.euler_route) turns them into commands.

    battery(room, capacity, battery)   Part 2's pick-cells battery problem, with c_v = cleaned
    two_robots(room)                   Part 3's two robots, sharing allowed, in the least time

Each starts from the earlier part's answer, so HiGHS has something to beat, and with a time limit
says how far off its best can still be.
"""
import copy
import math
import time
from collections import deque

import route_lp
from branch_cut import euler_route
from dirt_common import parse as parse_dirt
from dirt_common import route_with_cleaning


class Program:
    """A HiGHS model built a column block and a row at a time."""

    def __init__(self, time_limit):
        import highspy
        import numpy as np
        self.hs, self.np = highspy, np
        self.h = h = highspy.Highs()
        h.setOptionValue('output_flag', False)
        h.setOptionValue('time_limit', float(time_limit))
        h.setOptionValue('mip_rel_gap', 0.0)
        self.low, self.high, self.whole, self.cost = [], [], [], []

    def columns(self, count, high, whole, cost=0.0):
        first = len(self.low)
        self.low += [0.0] * count
        self.high += [float(high)] * count
        self.whole += [whole] * count
        self.cost += [float(cost)] * count
        return first

    def finish_columns(self, relax):
        np, h = self.np, self.h
        h.addVars(len(self.low), np.array(self.low), np.array(self.high))
        h.changeColsCost(len(self.cost), np.arange(len(self.cost), dtype=np.int32), np.array(self.cost))
        whole = [j for j, w in enumerate(self.whole) if w]
        if whole and not relax:
            h.changeColsIntegrality(len(whole), np.array(whole, dtype=np.int32),
                                    np.array([self.hs.HighsVarType.kInteger] * len(whole)))

    def row(self, low, high, coef):
        np = self.np
        self.h.addRow(float(low), float(high), len(coef), np.array(list(coef), dtype=np.int32),
                      np.array(list(coef.values()), dtype=float))


class Walk:
    """One robot's walk from field s: columns x (m), z (n), w (n), y (n), flow (2m) and its rows.
    `reach[v]` is False for fields it can never get to in time; `most` bounds the fields it reaches
    besides s, which is the most the flow carries over a pair."""

    def __init__(self, p, model, s, reach, most):
        n, m = model.n, model.m
        self.model, self.s = model, s
        self.X = p.columns(m, 2, True)
        self.Z = p.columns(n, 1, True)
        self.W = p.columns(n, 4, True)
        self.Y = p.columns(n, 1, True)
        self.F = p.columns(2 * m, most, False)
        self.p, self.reach, self.most = p, reach, most

    def rows(self):
        p, model, s = self.p, self.model, self.s
        n, m = model.n, model.m
        X, Z, W, Y, F = self.X, self.Z, self.W, self.Y, self.F
        inf = p.np.inf
        p.row(1, 1, {Y + s: 1})
        p.row(1, 1, {Z + v: 1 for v in range(n)})
        for v in range(n):
            if not self.reach[v]:
                p.row(0, 0, {Y + v: 1})
            coef = {X + e: 1 for e in model.touching[v]}           # parity, closed from the end to s
            coef[Z + v] = -1 if v == s else 1
            coef[W + v] = -2
            p.row(-1 if v == s else 0, -1 if v == s else 0, coef)
            for e in model.touching[v]:
                p.row(-inf, 0, {X + e: 1, Y + v: -2})
            if v != s:
                p.row(-inf, 0, {Z + v: 1, Y + v: -1})
                coef = {X + e: 1 for e in model.touching[v]}
                coef[Z + v] = 1
                coef[Y + v] = -2
                p.row(0, inf, coef)
            coef = {}                                               # flow: in - out = y_v
            for e in model.touching[v]:
                out, back = (F + 2 * e, F + 2 * e + 1) if model.edges[e][0] == v else (F + 2 * e + 1, F + 2 * e)
                coef[out] = coef.get(out, 0) + 1
                coef[back] = coef.get(back, 0) - 1
            if v == s:
                for u in range(n):
                    if u != s:
                        coef[Y + u] = -1
            else:
                coef[Y + v] = 1
            p.row(0, 0, coef)
        for e in range(m):
            p.row(-inf, 0, {F + 2 * e: 1, F + 2 * e + 1: 1, X + e: -self.most})
        coef = {X + e: -1 for e in range(m)}
        for v in range(n):
            coef[Y + v] = 1
        p.row(-inf, 1, coef)

    def moves(self):
        return {self.X + e: 1 for e in range(self.model.m)}

    def start(self, commands):
        """Column values for a walk given as commands ('.' waits): x, z, w, y and a flow along a
        breadth-first tree of its walked pairs."""
        model, s = self.model, self.s
        n, m = model.n, model.m
        index = {f: i for i, f in enumerate(model.fields)}
        edge = {frozenset(e): k for k, e in enumerate(model.edges)}
        step = {'^': (-1, 0), 'v': (1, 0), '<': (0, -1), '>': (0, 1)}
        x, seen = [0] * m, {s}
        r, c = model.fields[s]
        for ch in commands:
            if ch not in step:
                continue
            dr, dc = step[ch]
            a, b = index[(r, c)], index[(r + dr, c + dc)]
            x[edge[frozenset((a, b))]] += 1
            r, c = r + dr, c + dc
            seen.add(b)
        x = [v if v <= 2 else 2 - v % 2 for v in x]                # walked more than twice: 1 or 2 keeps parity
        z = [0] * n
        z[index[(r, c)]] = 1
        w = [(sum(x[k] for k in model.touching[v]) + (1 - z[v] if v == s else z[v])) // 2 for v in range(n)]
        y = [1 if v in seen else 0 for v in range(n)]
        parent, order, queue = {s: None}, [s], deque([s])
        while queue:
            u = queue.popleft()
            for k in model.touching[u]:
                if x[k]:
                    a, b = model.edges[k]
                    v = b if a == u else a
                    if v not in parent:
                        parent[v] = (u, k)
                        order.append(v)
                        queue.append(v)
        below = [0 if v == s else 1 for v in range(n)]
        f = [0.0] * (2 * m)
        for v in reversed(order[1:]):
            u, k = parent[v]
            below[u] += below[v]
            f[2 * k if model.edges[k][0] == u else 2 * k + 1] = float(below[v])
        return {self.X: x, self.Z: z, self.W: w, self.Y: y, self.F: f}

    def commands(self, values):
        model = copy.copy(self.model)
        model.s = self.s
        x = [round(values[self.X + e]) for e in range(model.m)]
        z = [round(values[self.Z + v]) for v in range(model.n)]
        return euler_route(model, x, z)


def distances(start, cells):
    dist, queue = {start: 0}, deque([start])
    while queue:
        r, c = queue.popleft()
        for nb in ((r - 1, c), (r + 1, c), (r, c - 1), (r, c + 1)):
            if nb in cells and nb not in dist:
                dist[nb] = dist[(r, c)] + 1
                queue.append(nb)
    return dist


def set_start(p, blocks, extra=()):
    """Hand HiGHS a first answer: blocks {first column: values}, extra [(column, value)]."""
    values = [0.0] * len(p.low)
    for first, vals in blocks.items():
        values[first:first + len(vals)] = [float(v) for v in vals]
    for col, v in extra:
        values[col] = float(v)
    sol = p.hs.HighsSolution()
    sol.col_value = values
    sol.value_valid = True
    p.h.setSolution(sol)


def battery(room, capacity, battery, time_limit=600, relax=False, start=None):
    """Part 2's battery problem, pick cells: the walk reaches fields, c_v says which it cleans.

        maximise sum dirt_v c_v   such that   c_v <= y_v,   sum dirt_v c_v <= capacity,   sum x <= battery

    `start` is a first route as commands (Part 2's heuristic by default). relax=True drops the
    whole-number rules and returns the linear program's value only. Otherwise returns a dict:
    commands, dirt (the best found), bound (no route collects more), optimal, nodes, seconds."""
    import pick_cells_battery
    t0 = time.perf_counter()
    dock, dirt = parse_dirt(room)
    model = route_lp.Model(room)
    n, s = model.n, model.s
    weight = [dirt[f] for f in model.fields]
    near = distances(dock, set(dirt))
    p = Program(time_limit)
    walk = Walk(p, model, s, [near[f] <= battery for f in model.fields], max(1, battery))
    C = p.columns(n, 1, True, cost=0)
    p.cost[C:C + n] = [-float(w) for w in weight]                  # HiGHS minimises: maximise the dirt
    p.finish_columns(relax)
    walk.rows()
    p.row(-p.np.inf, battery, walk.moves())
    p.row(-p.np.inf, capacity, {C + v: weight[v] for v in range(n)})
    for v in range(n):
        p.row(-p.np.inf, 0, {C + v: 1, walk.Y + v: -1})
    if relax:
        p.h.run()
        return -p.h.getInfo().objective_function_value
    first = start if start is not None else pick_cells_battery.solution(room, capacity, battery)
    blocks = walk.start(first)
    cleaned = cleaned_fields(dock, first)
    blocks[C] = [1 if f in cleaned else 0 for f in model.fields]
    set_start(p, blocks)
    p.h.run()
    info, status = p.h.getInfo(), p.h.getModelStatus()
    values = list(p.h.getSolution().col_value)
    moves = walk.commands(values)
    chosen = {model.fields[v] for v in range(n) if values[C + v] > 0.5}
    order, at, seen = [(dock, 0)], dock, {dock}
    step = {'^': (-1, 0), 'v': (1, 0), '<': (0, -1), '>': (0, 1)}
    for i, ch in enumerate(moves, 1):
        at = (at[0] + step[ch][0], at[1] + step[ch][1])
        if at not in seen:
            seen.add(at)
            order.append((at, i))
    commands = route_with_cleaning(list(moves), order, chosen)
    got = sum(dirt[f] for f in chosen)
    optimal = status == p.hs.HighsModelStatus.kOptimal
    bound = got if optimal else math.floor(-info.mip_dual_bound + 1e-6)
    return {'commands': commands, 'dirt': got, 'bound': bound, 'optimal': optimal,
            'nodes': info.mip_node_count, 'seconds': time.perf_counter() - t0}


def cleaned_fields(start, commands):
    step = {'^': (-1, 0), 'v': (1, 0), '<': (0, -1), '>': (0, 1)}
    at, out = start, set()
    for ch in commands:
        if ch == 'C':
            out.add(at)
        else:
            at = (at[0] + step[ch][0], at[1] + step[ch][1])
    return out


def two_robots(room, time_limit=600, relax=False, start=None):
    """Part 3's two robots, sharing allowed: two walks, every field reached by one of them, and
    the time T, the longer walk, as small as possible.

        minimise T   such that   y_A(v) + y_B(v) >= 1,   T >= sum x_A,   T >= sum x_B,
                                  T >= distance(start_k, v) * y_k(v)

    `start` is a first plan (Part 3's two_robots.solution by default); its time U also caps what
    each robot can reach. Returns the linear program's value with relax=True; otherwise a dict:
    plan (A's and B's commands), time, bound, optimal, nodes, seconds."""
    import two_robots as part3
    t0 = time.perf_counter()
    a, b, floor = part3.parse_two(room)
    model = route_lp.Model([row.replace('@', '.') for row in room])
    n = model.n
    index = {f: i for i, f in enumerate(model.fields)}
    plan = start if start is not None else part3.solution(room)
    most = part3.finish_time(plan)[0]
    p = Program(time_limit)
    walks, dist = [], []
    for s in (a, b):
        d = distances(s, floor)
        dist.append([d[f] for f in model.fields])
        walks.append(Walk(p, model, index[s], [d[f] <= most for f in model.fields], max(1, most)))
    T = p.columns(1, most, True, cost=1)
    p.finish_columns(relax)
    for k, walk in enumerate(walks):
        walk.rows()
        coef = walk.moves()
        coef[T] = -1
        p.row(-p.np.inf, 0, coef)
        for v in range(n):
            if dist[k][v]:
                p.row(-p.np.inf, 0, {walk.Y + v: dist[k][v], T: -1})
    for v in range(n):
        p.row(1, p.np.inf, {walks[0].Y + v: 1, walks[1].Y + v: 1})
    if relax:
        p.h.run()
        return p.h.getInfo().objective_function_value
    blocks = {}
    for walk, commands in zip(walks, plan):
        blocks.update(walk.start(commands))
    set_start(p, blocks, [(T, most)])
    p.h.run()
    info, status = p.h.getInfo(), p.h.getModelStatus()
    values = list(p.h.getSolution().col_value)
    best = tuple(walk.commands(values) for walk in walks)
    finish = max(len(c) for c in best)
    optimal = status == p.hs.HighsModelStatus.kOptimal
    bound = finish if optimal else math.ceil(info.mip_dual_bound - 1e-6)
    return {'plan': best, 'time': finish, 'bound': bound, 'optimal': optimal,
            'nodes': info.mip_node_count, 'seconds': time.perf_counter() - t0}


def no_collisions(room, plan):
    """Part 3's step 2 on any plan: keep it if the robots never collide, else replan one robot
    round the other (both ways round) or fall back on the plan that always works."""
    import two_robots_no_collide as part3
    from two_robots import finish_time, parse_two
    a, b, floor = parse_two(room)
    if part3.collision_free(a, b, *plan):
        return plan
    candidates = [part3.last_resort(a, b, floor)]
    for keep_a in (True, False):
        replanned = part3.replan(a, b, floor, plan, keep_a)
        if replanned and part3.collision_free(a, b, *replanned):
            candidates.append(replanned)
    return min(candidates, key=finish_time)