"""Checks and a report for Part 11: the battery by a price on every move, and both of Part 11's
integer programs (the battery, and two robots).

Run from this folder: python test_lagrange.py
Needs highspy (pip install highspy) for the integer programs.
"""
import itertools
import random
import time
from fractions import Fraction

import column_gen
import lagrange
import pick_cells_battery
import two_robots
import two_robots_exact
import walk_ip
from dirt_common import parse
from test_dirt import oracle_battery, random_room, simulate
from test_two_robots import oracle_sharing, simulate_two, two_start_room

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


def tiny(rng, most=12):
    """A random dirty room with at most `most` floor fields."""
    while True:
        room = random_room(rng, rng.randint(3, 6), rng.randint(3, 6), rng.choice([0, 0.2, 0.4]),
                           rng.choice(['12345', '5', '24', '35']))
        if len(parse(room)[1]) <= most:
            return room


def all_trips(inst):
    """Every elementary trip: every order of every set of fields that fits in the bag."""
    yield ()
    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 staircase(room, capacity):
    """Part 2's brute-force answer for every battery, until more battery stops helping."""
    f = [oracle_battery(room, capacity, 0, True)]
    most = lagrange.Instance(room, capacity).moves(lagrange.unlimited(lagrange.Instance(room, capacity)))
    for b in range(1, most + 1):
        f.append(oracle_battery(room, capacity, b, True))
    return f


def hull(f, battery):
    """The least concave function above the points (b, f[b]), at `battery` (f stays flat after)."""
    if battery >= len(f) - 1:
        return Fraction(f[-1])
    best = Fraction(f[battery])
    for i in range(battery + 1):
        for j in range(battery, len(f)):
            if i < j:
                best = max(best, f[i] + Fraction(f[j] - f[i], j - i) * (battery - i))
    return best


def test_inner(rng, count=150):
    """best_trip() against every trip there is: with memory of every field the labels find the best
    trip exactly; with a memory of one field, never a worse one."""
    for _ in range(count):
        room = tiny(rng, most=9)
        _, dirt = parse(room)
        cap = rng.randint(1, sum(dirt.values()) + 2)
        inst = lagrange.Instance(room, cap)
        if inst.n > 7:
            continue
        lam = Fraction(rng.randint(1, 24), rng.choice([1, 2, 4]))
        want = max(inst.load(t) - lam * inst.moves(t) for t in all_trips(inst))
        value, trip, moves, load = lagrange.best_trip(inst, lam)
        assert value == want == load - lam * moves and inst.moves(trip) == moves and inst.load(trip) == load
        small = lagrange.Instance(room, cap, memory=1)
        assert lagrange.best_trip(small, lam)[0] >= want
    print(f'inner problem ........... ok ({count} tiny rooms and prices: the labels find the best trip there is; '
          f'remembering one field, never a worse one)')


def test_dual(rng, count=60):
    """The exact dual is the concave hull of the best dirt for each battery, from Part 2's brute
    force: what the price can prove, and no more."""
    gaps = pairs = 0
    for _ in range(count):
        room = tiny(rng, most=10)
        _, dirt = parse(room)
        cap = rng.randint(1, sum(dirt.values()) + 2)
        f = staircase(room, cap)
        inst = lagrange.Instance(room, cap)
        for battery in range(len(f) + 1):
            value, _, _ = lagrange.dual(inst, battery)
            assert value == hull(f, battery), (room, cap, battery, value, hull(f, battery))
            best = f[min(battery, len(f) - 1)]
            assert value >= best
            gaps += value >= best + 1
            pairs += 1
    print(f'dual .................... ok ({count} tiny rooms, every battery: the price\'s best bound is the concave '
          f'hull of the brute-force answers; it could not prove the best route on {gaps} of {pairs} batteries)')


def test_subgradient(rng, count=60):
    """Subgradient steps never go below the dual, and every route they report is real."""
    off = []
    for _ in range(count):
        room = tiny(rng, most=10)
        _, dirt = parse(room)
        cap = rng.randint(1, sum(dirt.values()) + 2)
        battery = rng.randint(1, 12)
        inst = lagrange.Instance(room, cap)
        start = lagrange.trip_of(inst, pick_cells_battery.solution(room, cap, battery))
        bound, trip, _ = lagrange.subgradient(inst, battery, start)
        exact, _, _ = lagrange.dual(inst, battery)
        best = oracle_battery(room, cap, battery, True)
        commands = column_gen.commands(inst, [(trip, True, 0)]) if trip else ''
        assert simulate(room, commands, cap, battery=battery, pick=True) == inst.load(trip) <= best
        assert bound >= exact - 1e-9
        off.append(bound - float(exact))
    print(f'subgradient ............. ok ({count} tiny rooms: never below the dual, every route real; '
          f'{sum(o < 1e-6 for o in off)} reached the dual, the rest {max(off):.2f} above it at worst)')


def test_battery_ip(rng, count=80):
    for _ in range(count):
        room = tiny(rng)
        _, dirt = parse(room)
        cap, battery = rng.randint(1, sum(dirt.values()) + 2), rng.randint(0, 14)
        r = walk_ip.battery(room, cap, battery, 60)
        best = oracle_battery(room, cap, battery, True)
        assert r['optimal'] and r['dirt'] == best == simulate(room, r['commands'], cap, battery=battery, pick=True)
        assert walk_ip.battery(room, cap, battery, relax=True) >= best - 1e-6
    print(f'battery program ......... ok ({count} tiny rooms: proven best, equal to Part 2\'s brute force, and the '
          f'simulator agrees)')


def test_two_ip(rng, count=60):
    done = 0
    while done < count:
        room = two_start_room(rng, rng.randint(3, 6), rng.randint(3, 6), rng.choice([0, 0.2, 0.4]))
        if room is None or len(two_robots.parse_two(room)[2]) > 12:
            continue
        r = walk_ip.two_robots(room, 60)
        share = simulate_two(room, *two_robots_exact.solution(room, False), False)
        assert r['optimal'] and r['time'] == share == oracle_sharing(room) == simulate_two(room, *r['plan'], False)
        apart = simulate_two(room, *two_robots_exact.solution(room, True), True)
        assert simulate_two(room, *walk_ip.no_collisions(room, r['plan']), True) >= apart
        assert walk_ip.two_robots(room, relax=True) <= share + 1e-6
        done += 1
    print(f'two-robot program ....... ok ({count} tiny rooms: proven best, equal to Part 3\'s exact search and its '
          f'independent referee; without collisions too, every plan replayed)')


def report_example():
    print(f'\nPart 2\'s running example, bag {EXAMPLE_CAPACITY}:')
    print('  battery  Part 2  best  price bound  price  steps  program\'s LP  subgradient (60 steps)')
    inst = lagrange.Instance(EXAMPLE, EXAMPLE_CAPACITY)
    for battery in range(0, 11):
        heur = pick_cells_battery.solution(EXAMPLE, EXAMPLE_CAPACITY, battery)
        part2 = simulate(EXAMPLE, heur, EXAMPLE_CAPACITY, battery=battery, pick=True)
        best = walk_ip.battery(EXAMPLE, EXAMPLE_CAPACITY, battery)['dirt']
        value, lam, steps = lagrange.dual(inst, battery)
        lp = walk_ip.battery(EXAMPLE, EXAMPLE_CAPACITY, battery, relax=True)
        sub, _, k = lagrange.subgradient(inst, battery, lagrange.trip_of(inst, heur))
        print(f'  {battery:>7}  {part2:>6}  {best:>4}  {str(value):>11}  {str(lam):>5}  {steps:>5}  {lp:>12.2f}  '
              f'{sub:.3f} after {k}')


def battery_rooms(rng, groups, memory, seconds, steps=30):
    """Random rooms, 30% furniture, dirt 1 to 5: Part 2's heuristic, both bounds and the best. The
    price's bound is the best of `steps` subgradient steps from what a move earned on Part 2's
    route: the two-line method starts at prices so low that the labels hardly dominate."""
    print(f'\nThe battery on random rooms, 30% furniture, dirt 1 to 5 (the price remembers {memory} fields, '
          f'{steps} subgradient steps):')
    print('  size   bag  battery  rooms  Part 2  LP bound  price bound  best  proven  price proves  '
          'Part 2 best on  price time  program time')
    for size, cap, battery, count in groups:
        rows = []
        for _ in range(count):
            room = random_room(rng, size, size, 0.3, '12345')
            heur = pick_cells_battery.solution(room, cap, battery)
            part2 = simulate(room, heur, cap, battery=battery, pick=True)
            inst = lagrange.Instance(room, cap, memory)
            t = time.perf_counter()
            price, _, _ = lagrange.subgradient(inst, battery, lagrange.trip_of(inst, heur), steps=steps)
            t_price = time.perf_counter() - t
            lp = walk_ip.battery(room, cap, battery, relax=True)
            r = walk_ip.battery(room, cap, battery, seconds)
            assert simulate(room, r['commands'], cap, battery=battery, pick=True) == r['dirt'] >= part2
            assert price >= r['dirt'] - 1e-6 and lp >= r['dirt'] - 1e-6
            rows.append((part2, lp, float(price), r['dirt'], r['optimal'], price < r['dirt'] + 1 - 1e-6,
                         part2 == r['dirt'], t_price, r['seconds']))
        mean = lambda i: sum(x[i] for x in rows) / count
        print(f'  {size}x{size}  {cap:>4}  {battery:>7}  {count:>5}  {mean(0):>6.1f}  {mean(1):>8.1f}  {mean(2):>11.1f}  '
              f'{mean(3):>4.1f}  {sum(x[4] for x in rows):>3}/{count}  {sum(x[5] for x in rows):>8}/{count}  '
              f'{sum(x[6] for x in rows):>11}/{count}  {mean(7):>9.1f}s  {mean(8):>11.1f}s')


def two_rooms(rng, groups, seconds):
    """Random rooms with two robots: Part 3's plans, its bound, the program's bound and the best."""
    print('\nTwo robots on random rooms (furniture 0, 15 or 30%):')
    print('  size   rooms  fields  Part 3 bound  LP bound  Part 3 sharing  best sharing  proven  '
          'Part 3 apart  best apart  apart = best sharing  time')
    for size, count in groups:
        rows = []
        while len(rows) < count:
            room = two_start_room(rng, size, size, rng.choice([0, 0.15, 0.3]))
            if room is None:
                continue
            plan = two_robots.solution(room)
            share = simulate_two(room, *plan, False)
            r = walk_ip.two_robots(room, seconds, start=plan)
            lp = walk_ip.two_robots(room, relax=True, start=plan)
            best = simulate_two(room, *r['plan'], False)
            apart3 = simulate_two(room, *walk_ip.no_collisions(room, plan), True)
            apart = simulate_two(room, *walk_ip.no_collisions(room, r['plan']), True)
            rows.append((len(two_robots.parse_two(room)[2]), two_robots.lower_bound(room), lp, share, best,
                         r['optimal'], apart3, apart, apart == r['bound'], r['seconds'], r['bound']))
        mean = lambda i: sum(x[i] for x in rows) / count
        print(f'  {size}x{size}  {count:>5}  {mean(0):>6.0f}  {mean(1):>12.1f}  {mean(2):>8.1f}  {mean(3):>14.1f}  '
              f'{mean(4):>12.1f}  {sum(x[5] for x in rows):>3}/{count}  {mean(6):>12.1f}  {mean(7):>10.1f}  '
              f'{sum(x[8] for x in rows):>13}/{count}  {mean(9):>4.0f}s')
        for x in rows:
            if not x[5]:
                print(f'      not proven: best {x[4]}, bound {x[10]}')


def main():
    rng = random.Random(11)
    test_inner(rng)
    test_dual(rng)
    test_subgradient(rng)
    test_battery_ip(rng)
    test_two_ip(rng)
    report_example()
    battery_rooms(rng, ((8, 40, 10, 6), (10, 60, 15, 6), (12, 80, 20, 4)), memory=8, seconds=300)
    two_rooms(rng, ((8, 6), (10, 6), (12, 4)), seconds=300)


if __name__ == '__main__':
    main()