Advent of Code Day 20 - tilemakerpro.py

from collections import defaultdict, deque
from math import sqrt
from tile import Tile


AdjacencyMap = dict[int, list[int]]
ID = int


class TileMakerPro:
    def __init__(self, tiles: dict[ID, Tile]):
        self.tiles: dict[ID, Tile] = tiles
        self.dims: int = int(sqrt(len(self.tiles)))
        self.adjacency_map: AdjacencyMap = self.make_adjacency_map()
        self.corners: tuple[ID, ...] = self.find_corners()
        self.start: ID = self.corners[0]
        self.arranged: set[ID] = set()
        self.unlinked: set[ID] = set([tile for tile in self.tiles])
        self.map_is_arranged: bool = False
        self.image: list[list[str]] = []

    def make_adjacency_map(self) -> AdjacencyMap:
        tile_ids = [tile for tile in self.tiles.keys()]
        adjacency_map: AdjacencyMap = defaultdict(list)
        for id1 in tile_ids:
            for id2 in tile_ids:
                if id1 != id2 and self.tiles[id1].is_adjacent(self.tiles[id2]):
                    adjacency_map[id1].append(id2)
        return adjacency_map

    def find_corners(self) -> tuple[ID, ...]:
        return tuple(
            [id for id in self.adjacency_map if len(self.adjacency_map[id]) == 2]
        )

    def make_connection(self, t1: ID, t2: ID) -> bool:
        n_rotations = 0
        while True:
            if self.tiles[t1].right_edge == self.tiles[t2].left_edge:
                self.tiles[t1].adjacent["r"] = t2
                self.tiles[t2].adjacent["l"] = t1
                return True
            elif self.tiles[t1].left_edge == self.tiles[t2].right_edge:
                self.tiles[t1].adjacent["l"] = t2
                self.tiles[t2].adjacent["r"] = t1
                return True
            elif self.tiles[t1].top_edge == self.tiles[t2].bottom_edge:
                self.tiles[t1].adjacent["t"] = t2
                self.tiles[t2].adjacent["b"] = t1
                return True
            elif self.tiles[t1].bottom_edge == self.tiles[t2].top_edge:
                self.tiles[t1].adjacent["b"] = t2
                self.tiles[t2].adjacent["t"] = t1
                return True
            self.tiles[t2].rotate()
            n_rotations += 1
            if n_rotations % 4 == 0:
                self.tiles[t2].flip()
            if n_rotations % 8 == 0:
                return False

    def make_first_connection(self, t1: ID, t2: ID):
        n_rotations = 0
        while True:
            if self.make_connection(t1, t2):
                return True
            self.tiles[t1].rotate()
            n_rotations += 1
            if n_rotations % 4 == 0:
                self.tiles[t1].flip()

    def arrange_map(self):
        queue: deque[ID] = deque([self.start])
        while queue:
            current: ID = queue.popleft()
            if self.arranged:
                for tile in self.adjacency_map[current]:
                    self.make_connection(current, tile)
                    if tile not in self.arranged:
                        queue.append(tile)
                    self.arranged.add(current)
            else:
                # set first connections
                self.arranged.add(current)
                self.make_first_connection(current, self.adjacency_map[current][0])
                queue.append(self.adjacency_map[current][0])
                self.make_connection(current, self.adjacency_map[current][1])
                queue.append(self.adjacency_map[current][1])
        self.map_is_arranged = True

    def find_top_left_tile(self) -> ID:
        for tile in self.tiles:
            if all(
                [
                    self.tiles[tile].adjacent["b"] is not None,
                    self.tiles[tile].adjacent["t"] is None,
                    self.tiles[tile].adjacent["r"] is not None,
                    self.tiles[tile].adjacent["l"] is None,
                ]
            ):
                return tile

    def consolidate_tiles(self) -> Tile:
        if self.map_is_arranged is False:
            raise ValueError(
                "Incomplete map arrangement.  Call 'self.arrange_map() first."
            )
        image_map = []
        top_left = self.find_top_left_tile()
        current_row, current_col = top_left, top_left
        while current_row or current_col:
            row = [[] for _ in range(len(self.tiles[top_left].inner_tile))]
            while current_col:  # add rows
                for i in range(len(row)):
                    row[i].extend(self.tiles[current_col].inner_tile[i])
                current_col = self.tiles[current_col].adjacent["r"]
            image_map.extend(row)
            current_row = self.tiles[current_row].adjacent["b"]
            current_col = current_row
        return Tile(0, image_map)