mdgen.py

#!/usr/bin/env python3

"""Creates 2d mazes.

Two types of dungeon are supported

* BacktrackMaze -- No loops, all rooms connected. All nodes are corridors.

* MazeDungeon -- May have loops. All rooms connected.
                 Some dead-end nodes may be rooms instead of corridors.

Example:

    MASK =
    [
    "---XXXXXXXXX-------",
    "--XXXXXXXXXXXXXX---",
    "-XXXXXXXXXXXXXXXX--",
    "-XXXXXXXXXXXXXXXX--",
    "XXXXXXX--XXXXXXXXX-",
    "XXXXXXX--XXXXXXXXX-",
    "-XXXXXXXXXXXXXXXXXX",
    "--XXXXXXXXXXXXXXXX-",
    "----XXXXXXXXXX-----",
    "------XXXXXX-------",
    ]

    d = MazeDungeon(mask=MASK)  # MazeDungeons can have loops.
    d.generate()
    print(d)


          O--O  O--O--O--O  O--O--O
          |     |     |     |  |
          |     |     |     |  |
       @  O--O--O  @  O--O--O  O--O--O  O--O--O
       |  |     |  |        |     |  |  |  |  |
       |  |     |  |        |     |  |  |  |  |
    O--O  O  @--@  O--O--O  @  O  O--O  O  O  O--O
    |  |  |     |  |  |  |     |           |  |
    |  |  |     |  |  |  |     |           |  |
    @  O  O  O  @  O  @  @  O--O--O--O--O--O  O--@
       |  |  |     |        |  |     |
       |  |  |     |        |  |     |
 O--O--O  O--O  O--O        O--O  @  O--O--O--@  O--O
 |  |        |  |  |        |     |        |     |
 |  |        |  |  |        |     |        |     |
 O  O  O--O--O--O  O        O  @  O--O--O--O--O--O  @
    |  |           |           |  |     |        |  |
    |  |           |           |  |     |        |  |
    O--O  O--O--@  @  @  @  O--O  O--O--O  @  @  O--O--O
       |     |        |  |  |  |  |     |  |  |     |
       |     |        |  |  |  |  |     |  |  |     |
       O--O--O--O  O--O--O--O  O--O  @  O--O--O--O--O
                |           |        |  |
                |           |        |  |
             O--O--O  O--O  O  O--O--O--O
                   |  |  |  |  |
                   |  |  |  |  |
                   O--O  O--O  O--O

NB: Double lines are 2 way exits. Will be replaced with
single character funny arrows later, which will compact map.

You can also get a dunpy-format listing for easy parsing.

>>> d = MazeDungeon(rows=5, cols=5)
>>> print(d)

 O--O--O--O--O
    |  |  |  |
    |  |  |  |
 O--O  O--O--O
 |        |
 |        |
 O--O--O  O--O
 |     |     |
 |     |     |
 O--O  O  @  O
 |  |     |  |
 |  |     |  |
 @  O--O--O--O

>>> print(d.dunpy())

['room 1 None', 'room 2 None', 'room 3 None', 'room 4 None', 'room 5 None',
'room 6 None', 'room 7 None', 'room 8 None', 'room 9 None', 'room 10 None',
'room 11 None', 'room 12 None', 'room 13 None', 'room 14 None', 'room 15 None',
'room 16 None', 'room 17 None', 'room 18 None', 'room 19 Room', 'room 20 None',
'room 21 Room', 'room 22 None', 'room 23 None', 'room 24 None', 'room 25 None',
'exit 1 e 2', 'exit 2 s 7', 'exit 2 w 1', 'exit 2 e 3', 'exit 3 s 8', '
exit 3 w 2', 'exit 3 e 4', 'exit 4 s 9', 'exit 4 w 3', 'exit 4 e 5',
'exit 5 s 10', 'exit 5 w 4', 'exit 6 s 11', 'exit 6 e 7', 'exit 7 w 6',
'exit 7 n 2', 'exit 8 e 9', 'exit 8 n 3', 'exit 9 s 14', 'exit 9 w 8',
'exit 9 e 10', 'exit 9 n 4', 'exit 10 w 9', 'exit 10 n 5', 'exit 11 s 16',
'exit 11 e 12', 'exit 11 n 6', 'exit 12 w 11', 'exit 12 e 13', 'exit 13 s 18',
'exit 13 w 12', 'exit 14 e 15', 'exit 14 n 9', 'exit 15 s 20', 'exit 15 w 14',
'exit 16 s 21', 'exit 16 e 17', 'exit 16 n 11', 'exit 17 s 22', 'exit 17 w 16',
'exit 18 n 13', 'exit 19 s 24', 'exit 20 s 25', 'exit 20 n 15', 'exit 21 n 16',
'exit 22 e 23', 'exit 22 n 17', 'exit 23 w 22', 'exit 23 e 24', 'exit 24 w 23',
'exit 24 e 25', 'exit 24 n 19', 'exit 25 w 24', 'exit 25 n 20',
'color Room #RRGGBB', 'color None #RRGGBB']

"""

from enum import Enum
import random


class Directions (Enum):
    n = 0
    e = 1
    s = 2
    w = 3
    ne = 4
    se = 5
    sw = 6
    nw = 7
    # u = 8
    # d = 9

reverse_dictionary = {
    Directions.n: Directions.s,
    Directions.e: Directions.w,
    Directions.w: Directions.e,
    Directions.s: Directions.n,
    Directions.nw: Directions.se,
    Directions.ne: Directions.sw,
    Directions.se: Directions.nw,
    Directions.sw: Directions.ne,
}

g_direc_names ={
    Directions.n: 'n',
    Directions.e: 'e',
    Directions.w: 'w',
    Directions.s: 's',
    Directions.nw: 'nw',
    Directions.ne: 'ne',
    Directions.se: 'se',
    Directions.sw: 'sw',
}

# Row and column offsets.
g_offsets = {
    Directions.n: (-1, 0),
    Directions.ne: (-1, 1),
    Directions.e: (0, 1),
    Directions.se: (1, 1),
    Directions.s: (1, 0),
    Directions.sw: (1, -1),
    Directions.w: (0, -1),
    Directions.nw: (-1, -1),
}


# Why is this actually a function?
def reverse_direct(x):
    return reverse_dictionary[x]


class Cell:
    """Grid cell.

    Cells store their X, Y co-ords, a list of exits (adjacency matrix)
    and a bitflag indicating if they have been masked out. See __init__."""

    id_count = 0

    class Neighbour:

        """A grid cell exit.

        A cell has a dict (key is direction) of these, holding data on the
        adjacent cells. Masked cells have no adjacency references and cannot
        be reached by normal means. The tag is optional, allowing for
        generators to mark special rooms."""

        def __init__(self, direction, destination):
            self.direction = direction
            self.destination = destination
            self.connected = False

    def __str__(self):
        return "Cell(r{},c{})".format(self.r, self.c)

    def __init__(self, row, col, tag=None):
        """Create a new cell."""
        Cell.id_count += 1
        self.r = row
        self.c = col
        self.exits = {}
        self.masked = False  # Deleted cell, allowing for non-regular grids.
        self.id_no = Cell.id_count
        self.tag = tag
        self.backup = {}  # copy deleted exits here.

    def open_exit(self, direction):
        """Opens the door between two cells. Bidirecitonal."""
        neighbour = self.exits[direction].destination
        self.exits[direction].connected = True
        neighbour.exits[reverse_direct(direction)].connected = True

    def close_exit(self, direction):
        """Closes the door between two cells. Bidirecitonal."""
        neighbour = self.exits[direction].destination
        self.exits[direction].connected = False
        neighbour.exits[reverse_direct(direction)].connected = False

    def _unlink_exit(self, k, twoway=True):
        """Remove a cell.

        Uncreates the adjacency list entries between two cells, effectively
        unlinking a cell from the grid. Do not use this to close cell doors.
        The exit is actually copied, so it can be put back later if needed."""

        dest = self.exits[k].destination
        dest.backup[reverse_direct(k)] = dest.exits[reverse_direct(k)]
        self.backup[k] = self.exits[k]
        del dest.exits[reverse_direct(k)]
        del self.exits[k]

    def _link_exit(self, k):
        """Undo unlinking exit. Used after unmasking a cell."""
        dest = self.backup[k].destination
        dest.exits[reverse_direct(k)] = dest.backup[reverse_direct(k)]
        self.exits[k] = self.backup[k]

    def has_exit(self, d):
        """True if the cell has the specified exit, and the door is open."""
        try:
            return self.exits[d].connected
        except KeyError:
            return False

    def count_exits(self):
        """Number of connected, open exits."""
        dirs = [Directions.n, Directions.s, Directions.e, Directions.w]
        return sum([self.has_exit(d) for d in dirs])

    def mask(self):
        """Masks the cell, unlinking it from adjacency."""
        self.masked = True
        all_e = list(self.exits)
        for e in all_e:
            self._unlink_exit(e)

    def unmask(self):
        """Masks the cell, unlinking it from adjacency."""
        self.masked = False
        all_e = list(self.backup)
        for e in all_e:
            self._link_exit(e)

class Maze:
    """Maze is a wrapper around a 2d grid. It can be extended by over-riding
    generate() to make a fancy maze."""

    def __init__(self, **kwargs):
        """Creates a new maze.

        Keyword arguments:
        rows -- integer Y size
        cols -- integer X size
        mask -- list of strings

        If mask is not specified, a rectangular grid of [rows][cols] is
        created. If there is a mask, the grid is created from a mask defined
        as a list of strings. A mask has an 'X' if the cell exists and '-' if
        it doesn't. As an example, the following mask defines a circular maze
        with a hole in the center. The mask must be a perfect rectangle.
        [
        "-XXX-",
        "XX-XX",
        "XX-XX",
        " XXX ",
        ]

        The grid is created with all cells linked (in adjacency list) but
        all doors closed. Maze.generate() then builds a maze by selectively
        opening doors.
        """
        self.mask = kwargs.get('mask', None)
        if self.mask:
            self.num_rows, self.num_cols = self._dimensions_from_mask(self.mask)
        else:
            self.num_rows = kwargs.get('rows', 9)
            self.num_cols = kwargs.get('cols', 9)

        # Create 2d grid of cells.
        self.grid = [[Cell(r,c) for c in range(self.num_cols)]
                     for r in range(self.num_rows)]

        # Setup the adjacency list in each cell so that each cell has a link
        # to its neighbours.
        for cell in self.all_cells():
            row, col = cell.r, cell.c
            for d in [Directions.n, Directions.e, Directions.s, Directions.w]:
                rm, cm = g_offsets[d]
                if row + rm not in range(self.num_rows):
                    continue
                if col + cm not in range(self.num_cols):
                    continue
                cell.exits[d] = Cell.Neighbour(d, self.grid[row + rm][col + cm])

        # Apply mask if there is one.
        if self.mask:
            for ri, row in enumerate(self.mask):
                for ci, col in enumerate(row):
                    if col.upper() != 'X':
                        self.grid[ri][ci].mask()
            if not self.is_connected():
                raise RuntimeError("Error in mask.")

    def _dimensions_from_mask(self, data):
        cols = max(len(x) for x in data)
        rows = len(data)
        return rows, cols

    def all_cells(self):
        """Generator returning all cells in the grid (if not masked)."""
        for row in self.grid:
            for cell in row:
                if not cell.masked:
                    yield cell

    def random_cell(self):
        """Return a random (and not masked) grid cell."""
        return random.choice([x for x in self.all_cells()])

    def is_connected(self):
        """Check all spaces on grid are reachable.

        Return True if all cells are connected. False if not."""
        unvisited = set()
        for c in self.all_cells():
            unvisited.add(c)
            current = c
        stack = [current]
        while stack:
            current = stack.pop()
            if current in unvisited:
                unvisited.remove(current)
            for k, v in current.exits.items():
                if v.destination in unvisited:
                    stack.append(v.destination)
        for c in self.all_cells():
            if c in unvisited:
                return False
        return True

    def generate(self):
        """Generate the maze from a (masked) blank grid.

        Override in subclasses for custom maze types.
        You should check Maze.is_connected() in case you have a broken mask."""
        pass

    def set_connectivity(self, fraction):
        """Adds extra exits at random.

        x -- Fractional value (0.0 <= x <= 1.0).

        Exits are added completely randomly until the set fraction of total
        possible exits exist. This can add loops. This is intentional.
        """
        target = round(self._count_exits() * fraction)
        current = self._count_open_exits()
        while current < target:
            cell = self.random_cell()
            picks = list(cell.exits.keys())
            if not picks:
                continue
            pick = random.choice(picks)
            cell.open_exit(pick)
            current += 2

    def __repr__(self):
        return self.__str__()

    # TODO Use funny arrow characters to compact the display.
    def __str__(self):
        """Ascii representation of the maze.

        One way exits are shown by a shorter line."""
        output = []
        for row in self.grid:
            rows = ['', '', '']
            for cell in row:
                ch = 'O' if not cell.tag else '@'
                gap = [[' ', ' ', ' '], [' ', ch, ' '], [' ', ' ', ' ']]
                if cell.masked:
                    gap[1][1]=' '
                if cell.has_exit(Directions.n):
                    gap[0][1] = '|'
                if cell.has_exit(Directions.e):
                    gap[1][2] = '-'
                if cell.has_exit(Directions.w):
                    gap[1][0] = '-'
                if cell.has_exit(Directions.s):
                    gap[2][1] = '|'
                if cell.has_exit(Directions.ne):
                    gap[0][2] = '/'
                if cell.has_exit(Directions.nw):
                    gap[0][0] = '\\'
                if cell.has_exit(Directions.sw):
                    gap[2][0] = '/'
                if cell.has_exit(Directions.se):
                    gap[2][2] = '\\'
                rows = [rows[n] + "".join(gap[n]) for n in range(3)]
            for r in rows:
                output.append(r)
        return "\n".join(output)

    def show(self):
        print(self)

    def dunpy(self):
        '''Output the dungeon in Dunpy-compatible format.'''
        output = []
        tags = set()
        for cell in self.all_cells():
            output.append('room {} {}'.format(cell.id_no, cell.tag).strip())
            tags.add(cell.tag)
        for cell in self.all_cells():
            for k, v in cell.exits.items():
                if cell.has_exit(k):
                    output.append('exit {} {} {}'.format(cell.id_no,
                        g_direc_names[k], v.destination.id_no))
        for x in list(tags):
            output.append("color {} #RRGGBB".format(x))
        return output

    def _count_exits(self):
        '''Return the total number of exits possible.'''
        return sum ([len(c.exits) for c in self.all_cells()])

    def _count_open_exits(self):
        '''Return the total number of exits opened.'''
        total = 0
        for cell in self.all_cells():
            total += sum([v.connected for v in cell.exits.values()])
        return total


class BinaryMaze(Maze):
    """Low quality maze for testing purposes.

    Although a poor design, the maze is guaranteed to be complete and
    solvable."""
    def generate(self):
        if not self.is_connected():
            raise RuntimeError("Maze grid not connected. Check mask?")
        for row in self.grid:
            for cell in row:
                dirs = []
                if Directions.e in cell.exits:
                    dirs.append(Directions.e)
                if Directions.n in cell.exits:
                    dirs.append(Directions.n)
                if dirs:
                    cell.open_exit(random.choice(dirs))


class BacktrackMaze(Maze):
    """Implements recursive backtracker algorithm from Wikipedia.

    https://en.wikipedia.org/wiki/Maze_generation_algorithm"""

    def generate(self):
        if not self.is_connected():
            raise RuntimeError("Maze grid not connected. Check mask?")
        unvisited = set()
        for c in self.all_cells():
            unvisited.add(c)
        initial = self.random_cell()
        current = initial
        unvisited.remove(current)
        stack = []
        while unvisited:
            visit_list = []
            for k, v in current.exits.items():
                if v.destination in unvisited:
                    visit_list.append(k)
            if visit_list:
                picked = random.choice(visit_list)
                stack.append(current)
                current.open_exit(picked)
                current = current.exits[picked].destination
                unvisited.remove(current)
            elif stack:
                current = stack.pop()


class MazeDungeon(BacktrackMaze):
    """A maze dungeon is a maze with special rooms at the dead-ends.

    A maze dungeon is a regular maze with the deletion of some dead-end
    corridors, and their replacement with special rooms. Special rooms are
    shown by changing the cell tag from None to Room allowing the caller
    to identify them."""

    def generate(self):

        # Call backtrackmaze to generate.
        super().generate()
        self.set_connectivity(0.75)

        # list all dead-ends
        dead_ends = [c for c in self.all_cells() if c.count_exits() == 1]

        num_specials = min(self.num_rows * self.num_cols // 9, len(dead_ends))
        ends = random.sample(dead_ends, num_specials)
        for cell in ends:
            cell.tag = 'Room'


if __name__ == '__main__':
    with open('../data/mask2.txt', 'r') as fp:
        mask = fp.readlines()
    mask = [x.strip() for x in mask if x.strip()]

    for i in range(5):
        d = MazeDungeon(mask=mask)
        # d = MazeDungeon(rows=10, cols=5)
        try:
            d.generate()
            # print(d.dunpy())
            d.show()
            break
        except RuntimeError as e:
            print(e)