py-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_HORIZONTALLY = 1
    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

    def __trackBoundaries(self, pos: Coordinate):
        if pos.x < self.minX:
            self.minX = pos.x

        if pos.y < self.minY:
            self.minY = pos.y

        if pos.x > self.maxX:
            self.maxX = pos.x

        if pos.y > self.maxY:
            self.maxY = pos.y

    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):
        if value in OFF_STATES and pos in self.__grid:
            del self.__grid[pos]
        elif 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 getValue(self, pos: Coordinate) -> any:
        if pos not in self.__grid:
            return self.__default
        else:
            return self.__grid[pos]

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

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

    def isCorner(self, pos: Coordinate) -> bool:
        return pos in [
            Coordinate(self.minX, self.minY),
            Coordinate(self.minX, self.maxY),
            Coordinate(self.maxX, self.minY),
            Coordinate(self.maxX, self.maxY),
        ]

    def isWithinBoundaries(self, pos: Coordinate) -> bool:
        return self.minX <= pos.x <= self.maxX and self.minY <= pos.y <= self.maxY

    def add(self, pos: Coordinate, value: Union[float, int] = 1):
        if pos in self.__grid:
            self.__grid[pos] += value
        else:
            self.__trackBoundaries(pos)
            self.__grid[pos] = self.__default + value

    def sub(self, pos: Coordinate, value: Union[float, int] = 1):
        if pos in self.__grid:
            self.__grid[pos] -= value
        else:
            self.__trackBoundaries(pos)
            self.__grid[pos] = self.__default - value

    def getActiveCells(self):
        return [i for i in 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):
            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 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)