2020/python/day20.py

import fileinput
from pprint import pprint
import itertools
from copy import deepcopy
from typing import List, Tuple, Optional, Dict

N = 10
SNEK_POSITIONS = (
    (0, 0),
    (1, 1),
    (4, 1),
    (5, 0),
    (6, 0),
    (7, 1),
    (10, 1),
    (11, 0),
    (12, 0),
    (13, 1),
    (16, 1),
    (17, 0),
    (18, 0),
    (18, -1),
    (19, 0),
)

Tile = List[List[bool]]
Position = Tuple[int, int]
Extremes = Tuple[int, int, int, int]


def parse() -> List[Tuple[int, Tile]]:
    current_tile = None
    tiles = []
    for line in fileinput.input():
        line = line.strip()
        if not line:
            continue

        if line.startswith("Tile "):
            if current_tile:
                assert len(current_tile) == N
                tiles.append((current_id, current_tile))
            current_id = int(line[5:-1])
            current_tile = []
        else:
            assert len(line) == N
            current_tile.append([c == "#" for c in line])

    assert len(current_tile) == N
    tiles.append((current_id, current_tile))
    return tiles


def aligns_right(left: Tile, right: Tile) -> bool:
    return all(left_row[-1] == right_row[0] for (left_row, right_row) in zip(left, right))


def aligns_bottom(top: Tile, bottom: Tile) -> bool:
    return top[-1] == bottom[0]


def aligns(a: Tile, b: Tile) -> Optional[Position]:
    if aligns_bottom(a, b):
        return 0, 1
    if aligns_bottom(b, a):
        return 0, -1
    if aligns_right(a, b):
        return 1, 0
    if aligns_right(b, a):
        return -1, 0
    return None


def rotate(tile: Tile) -> Tile:
    output = list(reversed(tile))
    for y in range(len(tile)):
        for x in range(y):
            output[y][x], output[x][y] = output[x][y], output[y][x]
    return output


def flip(a: Tile) -> Tile:
    return list(reversed(a))


def rotate_align(a: Tile, b: Tile) -> Optional[Tuple[Tile, Position]]:
    """
    Rotates and flips a and checks if it aligns for every possible orientation.
    """
    for _ in range(4):
        if pos := aligns(a, b):
            return b, pos
        bf = flip(b)

        if pos := aligns(a, bf):
            return bf, pos
        b = rotate(b)
    return None


def part1(
    tiles: List[Tuple[int, Tile]]
) -> Tuple[Dict[Position, Tuple[int, Tile]], Extremes]:
    tile_positions = {tiles[0][0]: (0, 0)}
    position_tiles = {(0, 0): tiles[0]}

    while len(tile_positions) != len(tiles):
        for a_id, _ in tiles:
            try:
                (a_x, a_y) = a_pos = tile_positions[a_id]
                _, a_tile = position_tiles[a_pos]
            except KeyError:
                continue

            for b_id, b_tile in tiles:
                if b_id in tile_positions or a_id == b_id:
                    continue

                aligned = rotate_align(a_tile, b_tile)
                if aligned is not None:
                    transformed, b_pos = aligned
                    dx, dy = b_pos
                    b_x = a_x + dx
                    b_y = a_y + dy
                    tile_positions[b_id] = (b_x, b_y)
                    position_tiles[(b_x, b_y)] = b_id, transformed

    min_y = min(y for (_, y) in tile_positions.values())
    max_y = max(y for (_, y) in tile_positions.values())
    min_x = min(x for (x, _) in tile_positions.values())
    max_x = max(x for (x, _) in tile_positions.values())

    bl, _ = position_tiles[(min_x, min_y)]
    br, _ = position_tiles[(max_x, min_y)]
    tl, _ = position_tiles[(min_x, max_y)]
    tr, _ = position_tiles[(max_x, max_y)]
    print("Part 1:", tl * tr * bl * br)

    return position_tiles, (min_x, max_x, min_y, max_y)


def is_snek(image: Tile, start_x: int, start_y: int) -> bool:
    for (x, y) in SNEK_POSITIONS:
        try:
            if not image[y + start_y][x + start_x]:
                return False
        except:
            return False
    return True


def find_sneks(b: Tile) -> List[Position]:
    sneks = []
    for (x, y) in itertools.product(range(8 * 12), repeat=2):
        if is_snek(b, x, y):
            sneks.append((x, y))
    return sneks


def rotate_find_sneks(b: Tile) -> Tuple[Tile, List[Position]]:
    for _ in range(4):
        if position := find_sneks(b):
            return b, position
        bf = flip(b)

        if position := find_sneks(bf):
            return bf, position
        b = rotate(b)
    raise RuntimeError("no sneks found")


def remove_snek(image: Tile, snek_position: Position) -> None:
    # :(
    (x, y) = snek_position
    for (x_, y_) in SNEK_POSITIONS:
        image[y + y_][x + x_] = False


def part2(
    position_tiles: Dict[Position, Tuple[int, Tile]],
    extremes: Tuple[int, int, int, int],
) -> None:
    min_x, max_x, min_y, max_y = extremes
    image = [[False] * 8 * (max_x - min_x + 1) for _ in range(8 * (max_y - min_y + 1))]
    for (x, y), (_, tile) in position_tiles.items():
        x_ = (x - min_x) * 8
        y_ = (y - min_y) * 8
        for i, row in enumerate(tile[1:-1]):
            for j, c in enumerate(row[1:-1]):
                image[i + y_][j + x_] = c

    rotated_image, snek_positions = rotate_find_sneks(image)
    for snek_position in snek_positions:
        remove_snek(rotated_image, snek_position)

    part2 = 0
    for row in rotated_image:
        for c in row:
            if c:
                part2 += 1

    print("Part 2:", part2)


def main() -> None:
    tiles = parse()
    position_tiles, extremes = part1(tiles)
    part2(position_tiles, extremes)


if __name__ == "__main__":
    main()