"""Checks for the original cleaning robot solutions (Part 1 of the blog post).

Run from this folder: python test_cleaning_robot.py
"""
import random

from cleaning_robot import solution, solution_iterative, solution_trimmed

STEP = {'^': (-1, 0), 'v': (1, 0), '<': (0, -1), '>': (0, 1)}
SOLUTIONS = (solution, solution_trimmed, solution_iterative)

# The five example rooms, with the commands `solution` returns for each.
EXAMPLES = [
    (['######', '#....#', '#....#', '#...*#', '######'],
     '^^<vv<^^<vv^^>vv>^^>vv'),
    (['#######', '#.....#', '#.....#', '#..*..#', '#.....#', '#.....#', '#######'],
     '^^<vvvv<^^^^vvvv>>^>^^^>vvvv<>^^^^<vvv<v<^^^^>vv'),
    (['#########', '#..#....#', '#.*..#..#', '#....#..#', '#########'],
     '^<vv>>^>^>>vv>^^vv<^^<<vv^<v<<^^>v'),
    (['#######', '#*#...#', '#.#.#.#', '#...#.#', '###...#', '#######'],
     'vv>>^^>>vvv<<>>^^^<<vv<<^^'),
    (['#########', '####....#', '#......##', '#.*..#..#', '##...####', '#########'],
     '^<v^>>vv<>>^^^>v>^><vv><^<^<vvv<^^<v'),
]
TINY = ['####', '#*.#', '#..#', '####']


def simulate(room, cmds):
    """Replay cmds and return the set of cleaned fields; fail on a wall or a broken rule."""
    assert len(cmds) <= 50_000, 'more than 50,000 commands'
    assert set(cmds) <= set(STEP), f'unexpected command in {cmds!r}'
    R, C = len(room), len(room[0])
    r, c = next((r, c) for r in range(R) for c in range(C) if room[r][c] == '*')
    cleaned = {(r, c)}
    for ch in cmds:
        dr, dc = STEP[ch]
        r, c = r + dr, c + dc
        assert room[r][c] != '#', f'walked into a wall at {(r, c)}'
        cleaned.add((r, c))
    return cleaned


def floor(room):
    return {(r, c) for r, row in enumerate(room) for c, ch in enumerate(row) if ch != '#'}


def random_room(rng, R, C, wall_p):
    """A random room whose floor is all reachable from the start."""
    grid = [['#'] * C for _ in range(R)]
    for r in range(1, R - 1):
        for c in range(1, C - 1):
            if rng.random() >= wall_p:
                grid[r][c] = '.'
    cells = [(r, c) for r in range(R) for c in range(C) if grid[r][c] == '.']
    if not cells:
        grid[1][1] = '.'
        cells = [(1, 1)]
    start = rng.choice(cells)
    reach, stack = {start}, [start]
    while stack:
        r, c = stack.pop()
        for dr, dc in STEP.values():
            n = (r + dr, c + dc)
            if grid[n[0]][n[1]] == '.' and n not in reach:
                reach.add(n)
                stack.append(n)
    for r, c in cells:
        if (r, c) not in reach:
            grid[r][c] = '#'
    grid[start[0]][start[1]] = '*'
    return [''.join(row) for row in grid]


def empty_room(R, C, start):
    grid = [list('#' * C)] + [list('#' + '.' * (C - 2) + '#') for _ in range(R - 2)] + [list('#' * C)]
    grid[start[0]][start[1]] = '*'
    return [''.join(row) for row in grid]


def check_examples():
    for i, (room, expected) in enumerate(EXAMPLES, 1):
        for solve in SOLUTIONS:
            assert simulate(room, solve(room)) == floor(room), f'{solve.__name__} on example {i}'
        assert solution(room) == expected, f'example {i} commands changed'
    print('example rooms ........... all five cleaned by all three solutions')


def check_known_answers():
    assert solution(TINY) == 'v>^v<^'
    assert solution_trimmed(TINY) == 'v>^'
    print('tiny room ............... matches the blog post')


def check_random_rooms(count=3000):
    rng = random.Random(1)
    for _ in range(count):
        room = random_room(rng, rng.randint(3, 40), rng.randint(3, 40), rng.choice([0, 0.1, 0.3, 0.45]))
        full = solution(room)
        assert simulate(room, full) == floor(room)
        assert len(full) == 2 * (len(floor(room)) - 1)
        assert solution_iterative(room) == full
        short = solution_trimmed(room)
        assert full.startswith(short) and simulate(room, short) == floor(room)
    print(f'random rooms ............ {count} rooms cleaned, each in exactly 2 * (fields - 1) commands')


def check_largest_rooms():
    for start in [(1, 1), (1, 38), (38, 1), (38, 38), (20, 20)]:
        room = empty_room(40, 40, start)
        for solve in SOLUTIONS:
            assert simulate(room, solve(room)) == floor(room), (solve.__name__, start)
        assert len(solution(room)) == 2886
        if start != (20, 20):
            assert len(solution_trimmed(room)) == 1443
    print('empty 40x40 rooms ....... 2,886 commands from every corner and the centre (1,443 trimmed from a corner)')


def main():
    check_examples()
    check_known_answers()
    check_random_rooms()
    check_largest_rooms()
    print('all checks passed')


if __name__ == '__main__':
    main()