"""Checks and a report for Part 10: emptying the bag (pick cells) by column generation.

Run from this folder: python test_column_gen.py
Needs highspy (pip install highspy) for the master program and the integer program.
"""
import itertools
import random
import time

import column_gen
import dock_exact
import pick_cells_dock
from dirt_common import parse
from dock_common import bfs_tree, bounds
from test_dirt import random_room
from test_dock import MIXES, oracle_pick, simulate_dock

EXAMPLE = ['#######', '#*34#5#', '#15215#', '#4#325#', '#######']      # Part 2's running example


def tiny(rng, most=9):
    """A random room with at most `most` dirty fields, and a capacity."""
    while True:
        room = random_room(rng, rng.randint(3, 6), rng.randint(3, 6), rng.choice([0, 0.2, 0.4]), rng.choice(MIXES))
        _, dirt = parse(room)
        k = sum(1 for f in dirt if dirt[f])
        if 1 <= k <= most:
            return room, rng.choice([5, 6, 7, 8, 10, 15, max(5, sum(dirt.values()) + 1)])


def all_trips(inst):
    """Every elementary trip: every order of every set of fields that fits in the bag."""
    for k in range(1, inst.n + 1):
        for fields in itertools.combinations(range(inst.n), k):
            if sum(inst.dirt[j] for j in fields) <= inst.capacity:
                yield from itertools.permutations(fields)


def reduced(inst, trip, open_, pi, sigma):
    return inst.cost(trip, open_) - sum(pi[j] for j in trip) - (sigma if open_ else 0)


def test_pricing(rng, count=150):
    """price() against every trip there is. With a neighbourhood as big as the room, labels are
    elementary and the best trip must match exactly; with a small one, pricing may find a
    cheaper ng-route, never a dearer one. Random prices, including big ones."""
    exact = relaxed = 0
    for _ in range(count):
        room, cap = tiny(rng, most=7)
        inst = column_gen.Instance(room, cap)
        pi = [rng.choice([0, rng.randint(0, 12)]) for _ in range(inst.n)]
        sigma = rng.randint(-5, 10)
        best = min(min(reduced(inst, t, o, pi, sigma) for o in (False, True)) for t in all_trips(inst))
        found = column_gen.price(inst, pi, sigma)
        got = found[0][0] if found else 0.0
        assert abs(got - min(best, 0)) < 1e-9 or (got > -1e-9 and best > -1e-9), (room, cap, got, best)
        exact += 1
        small = column_gen.Instance(room, cap)
        small.ng = [(1 << j) | (1 << ((j + 1) % small.n)) for j in range(small.n)]     # neighbourhoods of two
        found = column_gen.price(small, pi, sigma)
        assert (found[0][0] if found else 0.0) <= min(best, 0) + 1e-9
        relaxed += 1
    print(f'pricing ................. ok ({exact} random prices on tiny rooms: the labels find the best trip '
          f'there is; with neighbourhoods of two, never a dearer one ({relaxed}))')


def test_completion_and_enumeration(rng, count=60):
    """completion() never overestimates the cheapest finish, and enumerate_trips() finds every
    trip with reduced cost within the gap (one order per set of fields)."""
    for _ in range(count):
        room, cap = tiny(rng, most=7)
        inst = column_gen.Instance(room, cap)
        pi = [rng.randint(0, 10) for _ in range(inst.n)]
        sigma = rng.randint(-3, 8)
        g = column_gen.completion(inst, pi, sigma)
        trips = list(all_trips(inst))
        for t in trips:                                 # every suffix of every trip
            for cut in range(1, len(t) + 1):
                head = t[:cut]
                load = sum(inst.dirt[j] for j in head)
                last = head[-1]
                rest = [x for x in trips if x[:len(head)] == head]
                finish = min(min(reduced(inst, x, o, pi, sigma) for o in (False, True)) - reduced_head(inst, head, pi) for x in rest)
                assert g[cap - load][last] <= finish + 1e-9
        gap = rng.randint(0, 6)
        want = {(frozenset(t), o) for t in trips for o in (False, True) if reduced(inst, t, o, pi, sigma) <= gap + 1e-9}
        got = {(frozenset(t), o) for t, o in column_gen.enumerate_trips(inst, pi, sigma, gap)}
        assert got == want, (room, cap)
    print(f'completion, enumeration . ok ({count} tiny rooms: the completion bound never overestimates, and '
          f'enumeration finds every trip within the gap)')


def reduced_head(inst, head, pi):
    n, d = inst.n, inst.dist
    return d[n][head[0]] + sum(d[a][b] for a, b in zip(head, head[1:])) - sum(pi[j] for j in head)


def test_masters(rng, count=20):
    """The exact master (simplex.py) and HiGHS reach the same bound."""
    for _ in range(count):
        room, cap = tiny(rng, most=7)
        exact = column_gen.generate(room, cap, solver='exact')[0]
        highs = column_gen.generate(room, cap)[0]
        assert abs(float(exact) - highs) < 1e-6, (room, cap, exact, highs)
    print(f'masters ................. ok ({count} tiny rooms: the exact simplex and HiGHS reach the same bound)')


def test_tiny_rooms(rng, count=120):
    """The whole method against Part 3's trip oracle and exact search."""
    below = 0
    for _ in range(count):
        room, cap = tiny(rng)
        r = column_gen.solve(room, cap)
        best = oracle_pick(room, cap)
        assert r['proven'] and r['moves'] == best == simulate_dock(room, r['commands'], cap, True)
        assert simulate_dock(room, dock_exact.solution(room, cap, True), cap, True) == best
        below += r['bound'] < best - 1e-6
    print(f'tiny rooms .............. ok ({count} rooms: proven best, equal to Part 3\'s trip oracle and exact '
          f'search, and the simulator agrees; the bound was below the best on {below})')


def report_example():
    print('\nPart 2\'s running example, pick cells:')
    print('  capacity  Part 3 fast  Part 3 lower  LP bound  best  proven  trips enumerated')
    for cap in (5, 10, 23):
        r = column_gen.solve(EXAMPLE, cap)
        dock, dirt = parse(EXAMPLE)
        lower, _ = bounds(dirt, bfs_tree(dock, dirt)[0], cap)
        fast = simulate_dock(EXAMPLE, pick_cells_dock.solution(EXAMPLE, cap), cap, True)
        print(f'  {cap:>8}  {fast:>11}  {lower:>12}  {r["bound"]:>8.2f}  {r["moves"]:>4}  {str(r["proven"]):>6}  '
              f'{r["enumerated"]:>16}')


def report_rooms(rng, groups=((8, 10, 6), (10, 5, 6), (10, 10, 6), (12, 5, 4), (12, 10, 4), (12, 20, 2))):
    print('\nRandom rooms, 30% furniture, dirt 1 to 5:')
    print('  size   capacity  rooms  fields  Part 3 fast  Part 3 lower  LP bound  best  proven  '
          'fast over best  time')
    for size, cap, count in groups:
        rows = []
        for _ in range(count):
            room = random_room(rng, size, size, 0.3, '12345')
            dock, dirt = parse(room)
            lower, _ = bounds(dirt, bfs_tree(dock, dirt)[0], cap)
            fast = simulate_dock(room, pick_cells_dock.solution(room, cap), cap, True)
            t = time.perf_counter()
            r = column_gen.solve(room, cap, time_limit=300)
            assert simulate_dock(room, r['commands'], cap, True) == r['moves']
            rows.append((sum(1 for f in dirt if dirt[f]), fast, lower, r['bound'], r['moves'], r['proven'],
                         time.perf_counter() - t))
        mean = lambda i: sum(r[i] for r in rows) / count
        print(f'  {size}x{size}  {cap:>8}  {count:>5}  {mean(0):>6.0f}  {mean(1):>11.1f}  {mean(2):>12.1f}  '
              f'{mean(3):>8.1f}  {mean(4):>4.1f}  {sum(r[5] for r in rows):>3}/{count}  '
              f'{mean(1) - mean(4):>14.1f}  {mean(6):>4.0f}s')
        for r in rows:
            if not r[5]:
                print(f'      not proven: best {r[4]}, bound {r[3]:.2f} ({100 * (r[4] - r[3]) / r[4]:.1f}% gap)')


def main():
    rng = random.Random(10)
    test_pricing(rng)
    test_completion_and_enumeration(rng)
    test_masters(rng)
    test_tiny_rooms(rng)
    report_example()
    report_rooms(rng)


if __name__ == '__main__':
    main()