py-tools/tools/grid.py

from __future__ import annotations
from .coordinate import Coordinate
from enum import Enum
from typing import Union, Any

OFF_STATES = [False, 0, None]
OFF = False
ON = True


class GridTransformation(Enum):
    FLIP_X = 1
    FLIP_HORIZONTALLY = 1  # alias for FLIP_X; prep for 3d-transformations
    FLIP_VERTICALLY = 2
    FLIP_DIAGONALLY = 3
    FLIP_DIAGONALLY_REV = 4
    ROTATE_RIGHT = 5
    ROTATE_LEFT = 6
    ROTATE_TWICE = 7


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

    def __trackBoundaries(self, pos: Coordinate):
        self.minX = min(self.minX, pos.x)
        self.minY = min(self.minY, pos.y)
        self.maxX = max(self.maxX, pos.x)
        self.maxY = max(self.maxY, pos.y)
        if self.mode3D:
            self.minZ = min(self.minZ, pos.z)
            self.maxZ = max(self.maxZ, pos.z)

    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 set(self, pos: Coordinate, value: Any = True):
        if pos.z is not None:
            self.mode3D = True

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

    def get(self, pos: Coordinate) -> Any:
        if pos in self.__grid:
            return self.__grid[pos]
        else:
            return self.__default

    def getState(self, pos: Coordinate) -> bool:
        if pos not in self.__grid:
            return False
        else:
            return self.__grid[pos] not in OFF_STATES

    def getOnCount(self):
        return len(self.__grid)

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

    def getCorners(self):
        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) -> bool:
        if self.mode3D:
            return self.minX <= pos.x <= self.maxX and self.minY <= pos.y <= self.maxY \
                   and self.minZ <= pos.z <= self.maxZ
        else:
            return self.minX <= pos.x <= self.maxX and self.minY <= pos.y <= self.maxY

    def add(self, pos: Coordinate, value: Union[float, int] = 1):
        self.set(pos, self.get(pos) + value)

    def sub(self, pos: Coordinate, value: Union[float, int] = 1):
        self.set(pos, self.get(pos) - value)

    def getActiveCells(self):
        return list(self.__grid.keys())

    def getSum(self, includeNegative: bool = True):
        grid_sum = 0
        for value in self.__grid.values():
            if includeNegative or value > 0:
                grid_sum += value

        return grid_sum

    def getNeighbourSum(self, pos: Coordinate, includeNegative: bool = True, includeDiagonal: bool = True) \
            -> Union[float, int]:
        neighbour_sum = 0
        for neighbour in pos.getNeighbours(
                includeDiagonal=includeDiagonal,
                minX=self.minX, minY=self.minY,
                maxX=self.maxX, maxY=self.maxY,
                minZ=self.minZ, maxZ=self.maxZ):
            if neighbour in self.__grid:
                if includeNegative or self.__grid[neighbour] > 0:
                    neighbour_sum += self.__grid[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 self.mode3D:
            raise NotImplementedError()  # that will take some time and thought

        if mode == GridTransformation.FLIP_HORIZONTALLY:
            for x in range(self.minX, (self.maxX - self.minX) // 2 + 1):
                for y in range(self.minY, self.maxY + 1):
                    self.flip(Coordinate(x, y), Coordinate(self.maxX - x, y))
        elif mode == GridTransformation.FLIP_VERTICALLY:
            for y in range(self.minY, (self.maxY - self.minY) // 2 + 1):
                for x in range(self.minX, self.maxX + 1):
                    self.flip(Coordinate(x, y), Coordinate(x, self.maxY - y))
        elif mode == GridTransformation.FLIP_DIAGONALLY:
            self.transform(GridTransformation.ROTATE_LEFT)
            self.transform(GridTransformation.FLIP_HORIZONTALLY)
        elif mode == GridTransformation.FLIP_DIAGONALLY_REV:
            self.transform(GridTransformation.ROTATE_RIGHT)
            self.transform(GridTransformation.FLIP_HORIZONTALLY)
        elif mode == GridTransformation.ROTATE_LEFT:
            newGrid = Grid()
            for x in range(self.maxX, self.minX - 1, -1):
                for y in range(self.minY, self.maxY + 1):
                    newGrid.set(Coordinate(y, self.maxX - x), self.get(Coordinate(x, y)))

            self.__dict__.update(newGrid.__dict__)
        elif mode == GridTransformation.ROTATE_RIGHT:
            newGrid = Grid()
            for x in range(self.minX, self.maxX + 1):
                for y in range(self.maxY, self.minY - 1, -1):
                    newGrid.set(Coordinate(self.maxY - y, x), self.get(Coordinate(x, y)))

            self.__dict__.update(newGrid.__dict__)
        elif mode == GridTransformation.ROTATE_TWICE:
            self.transform(GridTransformation.ROTATE_RIGHT)
            self.transform(GridTransformation.ROTATE_RIGHT)