Generates png maze with solution. CLI args -s SIZE of maze. (DEV)

maze-gen.py

import png
import sys
import time
import math
import random
import argparse
import svgwrite


class Point2d:
    def __init__(self, x=0, y=0):
        self.x = x
        self.y = y

    def __add__(self, pt):
        return Point2d(self.x + pt.x, self.y + pt.y)

    def __sub__(self, pt):
        return Point2d(self.x - pt.x, self.y - pt.y)

    def __abs__(self, pt):
        return math.sqrt((self.x - pt.x) ** 2 + (self.y - pt.y) ** 2)

    def __eq__(self, pt):
        return self.x == pt.x and self.y == pt.y

    def __repr__(self):
        return f"<Point2d [{self.x}, {self.y}]>"

    def __str__(self):
        return f"[{self.x}, {self.y}]"


class Maze:
    def __init__(self, size, start=Point2d(0, 0), end=Point2d(-1, -1)):
        self.size = size
        self.start = start
        self.pos = start
        self.end = end
        # For showing current progress
        self.path = [start]
        # For conversion to image
        self.movements = []
        self.visited = [start]

        if end.x == -1 and end.y == -1:
            self.end = Point2d(self.size.x - 1, self.size.y - 1)

    DIRECTIONS = [
        Point2d(1, 0),
        Point2d(-1, 0),
        Point2d(0, 1),
        Point2d(0, -1)
    ]

    def is_inside(self, pt):
        return 0 <= pt.x < self.size.x and 0 <= pt.y < self.size.y

    def get_unvisited_neighbours(self):
        neighs = []

        for direction in Maze.DIRECTIONS:
            tmp = self.pos + direction
            if self.is_inside(tmp) and tmp not in self.visited:
                neighs.append(tmp)

        return neighs

    def make_maze(self):
        while self.pos != self.end:
            moves = self.get_unvisited_neighbours()

            if len(moves) == 0:
                self.path.pop()
                self.pos = self.path[-1]
                continue

            self.pos = random.choice(moves)
            self.path.append(self.pos)
            self.visited.append(self.pos)
            self.movements.append([self.path[-2], self.path[-1]])

    def print_maze(self, char="O"):
        """Prints just the path to terminal, doesn't check terminal size."""
        sys.stdout.write("\x1b[2J")
        for node in self.path:
            sys.stdout.write(f"\x1b[{node.y + 1};{node.x + 1}H{char}")
        sys.stdout.flush()

    def maze_to_txt(self, solution=0):
        """Return 2D list containing data for image consumption.
        solution is value to be replaced for path coods except 0
        """
        # In an n x n grid there are n^2 squares and (n + 1)^2 lines
        l1 = [1 for x in range(self.size.x * 2 + 1)]
        l2 = [
            1 if x % 2 == 0 else 0
            for x in range(self.size.x * 2 + 1)
        ]
        grid = []

        for y in range(self.size.y):
            grid.append(l1.copy())
            grid.append(l2.copy())
        grid.append(l1.copy())

        # Remove walls from where it has moved
        # -1 <- and 1 -> movement
        for m in self.movements:
            start, end = m[0], m[1]
            vec = end - start
            step = Point2d(1, 1)

            if vec.x:
                step.x = 0 if vec.x > 0 else 2
            elif vec.y:
                step.y = 0 if vec.y > 0 else 2
            grid[2 * end.y + step.y][2 * end.x + step.x] = 0

        if solution:
            for node in self.path:
                grid[2 * node.y + 1][2 * node.x + 1] = solution

        return grid


def scale2d(grid, factor):
    """Scales up a 2D list by a given factor"""
    scaled = []

    for y in grid:
        tmp = []
        for x in y:
            tmp.extend([x for _t in range(factor)])
        scaled.extend([tmp for _t in range(factor)])

    return scaled


def main():
    parser = argparse.ArgumentParser("Maze generator")
    parser.add_argument("-s", "--size",  type=int, required=True)
    parser.add_argument("-S", "--square-size",  type=int, default=10)
    args = parser.parse_args()
    palette = [
        (255, 255, 255),
        (0, 0, 0),
        (255, 0, 0),
        (0, 255, 0)
    ]

    maze = Maze(Point2d(args.size, args.size))
    maze.make_maze()
    coverage = 100 * len(maze.visited) / args.size ** 2
    filename = random.randint(0, 2**32)

    print(f"Maze generated. Coverage: {coverage:.2f}%", )
    txt = maze.maze_to_txt()
    # For solution
    solution_txt = maze.maze_to_txt(solution=3)

    # Start and end markers red and green resptively
    txt[2 * maze.start.y + 1][2 * maze.start.y + 1] = 2
    txt[2 * maze.end.y + 1][2 * maze.end.y + 1] = 3

    # Scale up the grid
    print("Scaling up grid...")
    txt = scale2d(txt, args.square_size)
    solution_txt = scale2d(solution_txt, args.square_size)
    print(f"Writing png to {filename}.png and solution to {filename}.sol.png")

    with open(f"{filename}.png", "wb") as f:
        pngw = png.Writer(
            len(txt[0]), len(txt),
            bitdepth=2, palette=palette
        )
        pngw.write(f, txt)
    with open(f"{filename}.sol.png", "wb") as f:
        pngw = png.Writer(
            len(txt[0]), len(txt),
            bitdepth=2, palette=palette
        )
        pngw.write(f, solution_txt)

    print("Done.")


if __name__ == "__main__":
    main()