py-tools/src/tools/grid.py

from __future__ import annotations
import re
from .aoc_ocr import convert_array_6
from .coordinate import Coordinate, DistanceAlgorithm, Shape
from collections import deque
from collections.abc import Callable
from enum import Enum
from heapq import heappop, heappush
from math import inf
from typing import Any, Dict, List, Iterable, Mapping

OFF = False
ON = True


class GridTransformation(Enum):
    # Rotations always take the axis to rotate around as if it were the z-axis and then rotate clockwise
    # Counter-Rotations likewise, just anti-clockwise
    # 3D-only OPs have a number > 10
    ROTATE_Z = 3
    ROTATE_X = 11
    ROTATE_Y = 12
    COUNTER_ROTATE_X = 14
    COUNTER_ROTATE_Y = 15
    COUNTER_ROTATE_Z = 7
    FLIP_X = 4
    FLIP_Y = 5
    FLIP_Z = 13

    # Handy aliases
    FLIP_HORIZONTALLY = 5
    FLIP_VERTICALLY = 4
    ROTATE_RIGHT = 3
    ROTATE_LEFT = 7


class Grid:
    def __init__(self, default=False):
        self.__default = default
        self.__grid = {}
        self.minX = None
        self.minY = None
        self.maxX = None
        self.maxY = None
        self.minZ = None
        self.maxZ = None
        self.mode3D = False

    def __trackBoundaries(self, pos: Coordinate):
        if self.minX is None:
            self.minX, self.maxX, self.minY, self.maxY = pos[0], pos[0], pos[1], pos[1]
        else:
            self.minX = pos[0] if pos[0] < self.minX else self.minX
            self.minY = pos[1] if pos[1] < self.minY else self.minY
            self.maxX = pos[0] if pos[0] > self.maxX else self.maxX
            self.maxY = pos[1] if pos[1] > self.maxY else self.maxY

        if self.mode3D:
            if self.minZ is None:
                self.minZ = self.maxZ = pos[2]
            else:
                self.minZ = pos[2] if pos[2] < self.minZ else self.minZ
                self.maxZ = pos[2] if pos[2] > self.maxZ else self.maxZ

    def recalcBoundaries(self) -> None:
        self.minX, self.maxX, self.minY, self.maxY, self.minZ, self.maxZ = (
            None,
            None,
            None,
            None,
            None,
            None,
        )
        for c in self.__grid:
            self.__trackBoundaries(c)

    def getBoundaries(self) -> (int, int, int, int, int, int):
        if self.mode3D:
            return self.minX, self.minY, self.maxX, self.maxY, self.minZ, self.maxZ
        else:
            return self.minX, self.minY, self.maxX, self.maxY, -inf, inf

    def rangeX(self, pad: int = 0, reverse=False):
        if reverse:
            return range(self.maxX + pad, self.minX - pad - 1, -1)
        else:
            return range(self.minX - pad, self.maxX + pad + 1)

    def rangeY(self, pad: int = 0, reverse=False):
        if reverse:
            return range(self.maxY + pad, self.minY - pad - 1, -1)
        else:
            return range(self.minY - pad, self.maxY + pad + 1)

    def rangeZ(self, pad: int = 0, reverse=False):
        if not self.mode3D:
            raise ValueError("rangeZ not available in 2D space")
        if reverse:
            return range(self.maxZ + pad, self.minZ - pad - 1, -1)
        else:
            return range(self.minZ - pad, self.maxZ + pad + 1)

    def get_column(self, column: int) -> list[Any]:
        return [self.get(Coordinate(column, y)) for y in self.rangeY()]

    def get_row(self, row: int) -> list[Any]:
        return [self.get(Coordinate(x, row)) for x in self.rangeX()]

    def toggle(self, pos: Coordinate):
        if pos in self.__grid:
            del self.__grid[pos]
        else:
            self.__trackBoundaries(pos)
            self.__grid[pos] = not self.__default

    def toggleGrid(self):
        for x in self.rangeX():
            for y in self.rangeY():
                if not self.mode3D:
                    self.toggle(Coordinate(x, y))
                else:
                    for z in self.rangeZ():
                        self.toggle(Coordinate(x, y, z))

    def set(self, pos: Coordinate, value: Any = True) -> Any:
        if pos[2] is not None:
            self.mode3D = True

        if (value == self.__default) and pos in self.__grid:
            del self.__grid[pos]
        elif value != self.__default:
            self.__trackBoundaries(pos)
            self.__grid[pos] = value

        return value

    def move(
        self,
        pos: Coordinate,
        vec: Coordinate,
    ):
        target = pos + vec
        self.set(target, self.get(pos))
        if pos in self.__grid:
            del self.__grid[pos]

    def add(self, pos: Coordinate, value: int | float = 1) -> int | float:
        return self.set(pos, self.get(pos) + value)

    def sub(self, pos: Coordinate, value: int | float = 1) -> int | float:
        return self.set(pos, self.get(pos) - value)

    def mul(self, pos: Coordinate, value: int | float = 1) -> int | float:
        return self.set(pos, self.get(pos) * value)

    def div(self, pos: Coordinate, value: int | float = 1) -> int | float:
        return self.set(pos, self.get(pos) / value)

    def add_shape(self, shape: Shape, value: int | float = 1) -> None:
        for x in range(shape.top_left[0], shape.bottom_right[0] + 1):
            for y in range(shape.top_left[1], shape.bottom_right[1] + 1):
                if not shape.mode_3d:
                    pos = Coordinate(x, y)
                    self.set(pos, self.get(pos) + value)
                else:
                    for z in range(shape.top_left[2], shape.bottom_right[2] + 1):
                        pos = Coordinate(x, y, z)
                        self.set(pos, self.get(pos) + value)

    def get(self, pos: Coordinate) -> Any:
        return self.__grid.get(pos, self.__default)

    def getOnCount(self) -> int:
        return len(self.__grid)

    def count(self, value: Any) -> int:
        return list(self.__grid.values()).count(value)

    def isSet(self, pos: Coordinate) -> bool:
        return pos in self.__grid

    def getCorners(self) -> List[Coordinate]:
        if not self.mode3D:
            return [
                Coordinate(self.minX, self.minY),
                Coordinate(self.minX, self.maxY),
                Coordinate(self.maxX, self.minY),
                Coordinate(self.maxX, self.maxY),
            ]
        else:
            return [
                Coordinate(self.minX, self.minY, self.minZ),
                Coordinate(self.minX, self.minY, self.maxZ),
                Coordinate(self.minX, self.maxY, self.minZ),
                Coordinate(self.minX, self.maxY, self.maxZ),
                Coordinate(self.maxX, self.minY, self.minZ),
                Coordinate(self.maxX, self.minY, self.maxZ),
                Coordinate(self.maxX, self.maxY, self.minZ),
                Coordinate(self.maxX, self.maxY, self.maxZ),
            ]

    def isCorner(self, pos: Coordinate) -> bool:
        return pos in self.getCorners()

    def isWithinBoundaries(self, pos: Coordinate, pad: int = 0) -> bool:
        if self.mode3D:
            return (
                self.minX + pad <= pos[0] <= self.maxX - pad
                and self.minY + pad <= pos[1] <= self.maxY - pad
                and self.minZ + pad <= pos[2] <= self.maxZ - pad
            )
        else:
            return (
                self.minX + pad <= pos[0] <= self.maxX - pad
                and self.minY + pad <= pos[1] <= self.maxY - pad
            )

    def getActiveCells(
        self, x: int = None, y: int = None, z: int = None
    ) -> Iterable[Coordinate]:
        if x is not None or y is not None or z is not None:
            return (
                c
                for c in self.__grid.keys()
                if (c[0] == x if x is not None else True)
                and (c[1] == y if y is not None else True)
                and (c[2] == z if z is not None else True)
            )
        else:
            return self.__grid.keys()

    def getActiveRegion(
        self,
        start: Coordinate,
        includeDiagonal: bool = False,
        ignore: List[Coordinate] = None,
    ) -> List[Coordinate]:
        if not self.get(start):
            return []
        if ignore is None:
            ignore = []
        ignore.append(start)
        for c in self.getNeighboursOf(start, includeDiagonal=includeDiagonal):
            if c not in ignore:
                ignore = self.getActiveRegion(c, includeDiagonal, ignore)

        return ignore

    def values(self):
        return self.__grid.values()

    def getSum(self, includeNegative: bool = True) -> int | float:
        if not self.mode3D:
            return sum(
                self.get(Coordinate(x, y))
                for x in self.rangeX()
                for y in self.rangeY()
                if includeNegative or self.get(Coordinate(x, y)) >= 0
            )
        else:
            return sum(
                self.get(Coordinate(x, y, z))
                for x in self.rangeX()
                for y in self.rangeY()
                for z in self.rangeZ()
                if includeNegative or self.get(Coordinate(x, y)) >= 0
            )

    def getNeighboursOf(
        self,
        pos: Coordinate,
        includeDefault: bool = False,
        includeDiagonal: bool = True,
    ) -> Iterable[Coordinate]:
        neighbours = pos.getNeighbours(
            includeDiagonal=includeDiagonal,
            minX=self.minX,
            minY=self.minY,
            minZ=self.minZ,
            maxX=self.maxX,
            maxY=self.maxY,
            maxZ=self.maxZ,
        )
        for x in neighbours:
            if includeDefault or x in self.__grid:
                yield x

    def getNeighbourSum(
        self,
        pos: Coordinate,
        includeNegative: bool = True,
        includeDiagonal: bool = True,
    ) -> int | float:
        neighbour_sum = 0
        for neighbour in self.getNeighboursOf(pos, includeDefault=includeDiagonal):
            if includeNegative or self.get(neighbour) > 0:
                neighbour_sum += self.get(neighbour)

        return neighbour_sum

    def flip(self, c1: Coordinate, c2: Coordinate):
        buf = self.get(c1)
        self.set(c1, self.get(c2))
        self.set(c2, buf)

    def transform(self, mode: GridTransformation):
        if mode.value > 10 and not self.mode3D:
            raise ValueError("Operation not possible in 2D space", mode)

        coords = self.__grid
        self.__grid = {}
        if mode == GridTransformation.ROTATE_X:
            shift_z = self.maxY
            for c, v in coords.items():
                self.set(Coordinate(c[0], c[2], -c[1]), v)
            self.shift(shift_z=shift_z)
        elif mode == GridTransformation.ROTATE_Y:
            shift_x = self.maxX
            for c, v in coords.items():
                self.set(Coordinate(-c[2], c[1], c[0]), v)
            self.shift(shift_x=shift_x)
        elif mode == GridTransformation.ROTATE_Z:
            shift_x = self.maxX
            for c, v in coords.items():
                self.set(Coordinate(-c[1], c[0], c[2]), v)
            self.shift(shift_x=shift_x)
        elif mode == GridTransformation.COUNTER_ROTATE_X:
            shift_y = self.maxY
            for c, v in coords.items():
                self.set(Coordinate(c[0], -c[2], c[1]), v)
            self.shift(shift_y=shift_y)
        elif mode == GridTransformation.COUNTER_ROTATE_Y:
            shift_z = self.maxZ
            for c, v in coords.items():
                self.set(Coordinate(c[2], c[1], -c[0]), v)
            self.shift(shift_z=shift_z)
        elif mode == GridTransformation.COUNTER_ROTATE_Z:
            shift_y = self.maxY
            for c, v in coords.items():
                self.set(Coordinate(c[1], -c[0], c[2]), v)
            self.shift(shift_y=shift_y)
        elif mode == GridTransformation.FLIP_X:
            shift_x = self.maxX
            for c, v in coords.items():
                self.set(Coordinate(-c[0], c[1], c[2]), v)
            self.shift(shift_x=shift_x)
        elif mode == GridTransformation.FLIP_Y:
            shift_y = self.maxY
            for c, v in coords.items():
                self.set(Coordinate(c[0], -c[1], c[2]), v)
            self.shift(shift_y=shift_y)
        elif mode == GridTransformation.FLIP_Z:
            shift_z = self.maxZ
            for c, v in coords.items():
                self.set(Coordinate(c[0], c[1], -c[2]), v)
            self.shift(shift_z=shift_z)
        else:
            raise NotImplementedError(mode)

        self.recalcBoundaries()

    def shift(self, shift_x: int = 0, shift_y: int = 0, shift_z: int = 0):
        self.minX, self.minY = self.minX + shift_x, self.minY + shift_y
        self.maxX, self.maxY = self.maxX + shift_x, self.maxY + shift_y
        if self.mode3D:
            self.minZ, self.maxZ = self.minZ + shift_z, self.maxZ + shift_z
        coords = self.__grid
        self.__grid = {}
        for c, v in coords.items():
            if self.mode3D:
                nc = Coordinate(c[0] + shift_x, c[1] + shift_y, c[2] + shift_z)
            else:
                nc = Coordinate(c[0] + shift_x, c[1] + shift_y)
            self.set(nc, v)

    def shift_zero(self, recalc: bool = True):
        # self.shift() to (0, 0, 0) being top, left, front
        if recalc:
            self.recalcBoundaries()
        if self.mode3D:
            self.shift(0 - self.minX, 0 - self.minY, 0 - self.minZ)
        else:
            self.shift(0 - self.minX, 0 - self.minY)

    def getPath_BFS(
        self,
        pos_from: Coordinate,
        pos_to: Coordinate,
        includeDiagonal: bool,
        walls: List[Any] = None,
        stop_at_first: Any = None,
    ) -> List[Coordinate] | None:
        queue = deque()
        came_from = {pos_from: None}
        queue.append(pos_from)
        if walls is None:
            walls = [self.__default]

        while queue:
            current = queue.popleft()
            found_end = False
            for c in self.getNeighboursOf(
                current,
                includeDiagonal=includeDiagonal,
                includeDefault=self.__default not in walls,
            ):
                if c in came_from and self.get(c) in walls:
                    continue
                came_from[c] = current
                if c == pos_to or (
                    stop_at_first is not None and self.get(c) == stop_at_first
                ):
                    pos_to = c
                    found_end = True
                    break
                queue.append(c)
            if found_end:
                break

        if pos_to not in came_from:
            return None

        ret = []
        while pos_to in came_from:
            ret.insert(0, pos_to)
            pos_to = came_from[pos_to]

        return ret

    def getPath(
        self,
        pos_from: Coordinate,
        pos_to: Coordinate,
        includeDiagonal: bool,
        walls: List[Any] = None,
        weighted: bool = False,
    ) -> List[Coordinate] | None:
        f_costs = []
        if walls is None:
            walls = [self.__default]

        openNodes: Dict[Coordinate, tuple] = {}
        closedNodes: Dict[Coordinate, tuple] = {}

        openNodes[pos_from] = (0, pos_from.getDistanceTo(pos_to), None)
        heappush(f_costs, (0, pos_from))

        while f_costs:
            _, currentCoord = heappop(f_costs)
            if currentCoord not in openNodes:
                continue
            currentNode = openNodes[currentCoord]

            closedNodes[currentCoord] = currentNode
            del openNodes[currentCoord]
            if currentCoord == pos_to:
                break

            for neighbour in self.getNeighboursOf(
                currentCoord, includeDefault=True, includeDiagonal=includeDiagonal
            ):
                if self.get(neighbour) in walls or neighbour in closedNodes:
                    continue

                if weighted:
                    neighbourDist = self.get(neighbour)
                elif not includeDiagonal:
                    neighbourDist = 1
                else:
                    neighbourDist = currentCoord.getDistanceTo(
                        neighbour, DistanceAlgorithm.MANHATTAN, includeDiagonal
                    )

                targetDist = neighbour.getDistanceTo(pos_to)
                f_cost = targetDist + neighbourDist + currentNode[1]

                if neighbour not in openNodes or f_cost < openNodes[neighbour][0]:
                    openNodes[neighbour] = (
                        f_cost,
                        currentNode[1] + neighbourDist,
                        currentCoord,
                    )
                    heappush(f_costs, (f_cost, neighbour))

        if pos_to not in closedNodes:
            return None
        else:
            currentNode = closedNodes[pos_to]
            pathCoords = [pos_to]
            while currentNode[2]:
                pathCoords.append(currentNode[2])
                currentNode = closedNodes[currentNode[2]]

            return pathCoords

    def sub_grid(
        self,
        from_x: int,
        from_y: int,
        to_x: int,
        to_y: int,
        from_z: int = None,
        to_z: int = None,
    ) -> "Grid":
        if self.mode3D and (from_z is None or to_z is None):
            raise ValueError(
                "sub_grid() on mode3d Grids requires from_z and to_z to be set"
            )
        count_x, count_y, count_z = 0, 0, 0
        new_grid = Grid(self.__default)
        for x in range(from_x, to_x + 1):
            for y in range(from_y, to_y + 1):
                if not self.mode3D:
                    new_grid.set(
                        Coordinate(count_x, count_y), self.get(Coordinate(x, y))
                    )
                else:
                    for z in range(from_z, to_z + 1):
                        new_grid.set(
                            Coordinate(count_x, count_y, count_z),
                            self.get(Coordinate(x, y, z)),
                        )
                        count_z += 1

                    count_z = 0
                count_y += 1
            count_y = 0
            count_x += 1

        return new_grid

    def update(self, x: int, y: int, grid: Grid) -> None:
        put_x, put_y = x, y
        for get_x in grid.rangeX():
            for get_y in grid.rangeY():
                self.set(Coordinate(put_x, put_y), grid.get(Coordinate(get_x, get_y)))
                put_y += 1
            put_y = y
            put_x += 1

    def print(
        self,
        spacer: str = "",
        true_char: str = "#",
        false_char: str = " ",
        translate: dict = None,
        mark: list = None,
        z_level: int = None,
        bool_mode: bool = False,
    ):
        if translate is None:
            translate = {}

        if true_char is not None and True not in translate:
            translate[True] = true_char
        if false_char is not None and False not in translate:
            translate[False] = false_char

        for y in range(self.minY, self.maxY + 1):
            for x in range(self.minX, self.maxX + 1):
                pos = Coordinate(x, y, z_level)

                if mark and pos in mark:
                    print("X", end="")
                elif bool_mode:
                    print(true_char if self.get(pos) else false_char, end="")
                else:
                    value = self.get(pos)
                    if isinstance(value, list):
                        value = len(value)

                    if isinstance(value, Enum):
                        value = value.value

                    print(value if value not in translate else translate[value], end="")
                print(spacer, end="")

            print()

    def get_aoc_ocr_string(self, x_shift: int = 0, y_shift: int = 0):
        return convert_array_6(
            [
                [
                    "#" if self.get(Coordinate(x + x_shift, y + y_shift)) else "."
                    for x in self.rangeX()
                ]
                for y in self.rangeY()
            ]
        )

    def __str__(self, true_char: str = "#", false_char: str = "."):
        return "/".join(
            "".join(
                true_char if self.get(Coordinate(x, y)) else false_char
                for x in range(self.minX, self.maxX + 1)
            )
            for y in range(self.minY, self.maxY + 1)
        )

    @classmethod
    def from_str(
        cls,
        grid_string: str,
        default: Any = False,
        true_char: str = "#",
        true_value: Any = True,
        translate: dict = None,
        mode3d: bool = False,
    ) -> "Grid":
        if translate is None:
            translate = {}
        if (
            true_char is not None
            and True not in translate.values()
            and true_char not in translate
        ):
            translate[true_char] = true_value if true_value is not None else True

        ret = cls(default=default)
        for y, line in enumerate(grid_string.split("/")):
            for x, c in enumerate(line):
                if mode3d:
                    coord = Coordinate(x, y, 0)
                else:
                    coord = Coordinate(x, y)

                if c in translate:
                    ret.set(coord, translate[c])
                else:
                    ret.set(coord, c)

        return ret

    @classmethod
    def from_data(
        cls,
        data: Iterable[Iterable],
        default: Any = False,
        translate: Mapping[str, Any] = None,
        gen_3d: bool = False,
    ) -> Grid:
        """
        Every entry in data will be treated as row, every entry in data[entry] will be a separate column.
        gen_3d = True will just add z=0 to every Coordinate
        translate is used on every data[entry] and if present as key, its value will be used instead
            a value in translate can be a function with the following signature: def translate(value: Any) -> Any
            a key in translate is either a string of len 1 or it will be treated as regexp
                if multiple regexp match, the first encountered wins
                if there is a key that matches the entry it wins over any mathing regexp
        """
        grid = cls(default=default)

        regex_in_translate = False
        if translate is not None:
            for k in translate:
                if len(k) > 1:
                    regex_in_translate = True

        for y, row in enumerate(data):
            for x, col in enumerate(row):
                if translate is not None and col in translate:
                    if isinstance(translate[col], Callable):
                        col = translate[col](col)
                    else:
                        col = translate[col]
                elif regex_in_translate:
                    for k, v in translate.items():
                        if len(k) == 1:
                            continue

                        if re.search(k, col):
                            if isinstance(v, Callable):
                                col = translate[k](col)
                            else:
                                col = v
                            break

                if gen_3d:
                    grid.set(Coordinate(x, y, 0), col)
                else:
                    grid.set(Coordinate(x, y), col)

        return grid

    def __hash__(self):
        return hash(frozenset(self.__grid.items()))