squirrel-battle/squirrelbattle/mapgeneration/randomwalk.py

# Copyright (C) 2020 by ÿnérant, eichhornchen, nicomarg, charlse
# SPDX-License-Identifier: GPL-3.0-or-later

from enum import auto, Enum
from random import choice, random, randint
from typing import Tuple

from ..interfaces import Map, Tile


DEFAULT_PARAMS = {
    "split_chance": .15,
    "turn_chance": .5,
    "death_chance": .1,
    "max_walkers": 15,
    "width": 100,
    "height": 100,
    "fill": .4,
    "no_lone_walls": False,
}


class Directions(Enum):
    up = auto()
    down = auto()
    left = auto()
    right = auto()


class Walker:
    def __init__(self, x: int, y: int):
        self.x = x
        self.y = y
        self.dir = choice(list(Directions))

    def random_turn(self) -> None:
        self.dir = choice(list(Directions))

    def next_pos(self) -> Tuple[int, int]:
        if self.dir == Directions.up:
            return self.x, self.y + 1
        elif self.dir == Directions.down:
            return self.x, self.y - 1
        elif self.dir == Directions.right:
            return self.x + 1, self.y
        elif self.dir == Directions.left:
            return self.x - 1, self.y

    def move_in_bounds(self, width: int, height: int) -> None:
        nx, ny = self.next_pos()
        if 0 < nx < width and 0 < ny < height:
            self.x, self.y = nx, ny

    def split(self) -> "Walker":
        child = Walker(self.x, self.y)
        child.dir = self.dir
        return child


class Generator:
    def __init__(self, params: dict = DEFAULT_PARAMS):
        self.params = params

    def run(self) -> Map:
        width, height = self.params["width"], self.params["height"]
        walkers = [Walker(width // 2, height // 2)]
        grid = [[Tile.EMPTY for _ in range(width)] for _ in range(height)]
        count = 0
        while count < self.params["fill"] * width * height:
            # because we can't add or remove walkers while looping over the pop
            # we need lists to keep track of what will be the walkers for the
            # next iteration of the main loop
            next_walker_pop = []

            for walker in walkers:
                if grid[walker.y][walker.x] == Tile.EMPTY:
                    count += 1
                    grid[walker.y][walker.x] = Tile.FLOOR
                if random() < self.params["turn_chance"]:
                    walker.random_turn()
                walker.move_in_bounds(width, height)
                if random() > self.params["death_chance"]:
                    next_walker_pop.append(walker)

            # we make sure to never kill all walkers
            if not next_walker_pop:
                next_walker_pop.append(choice(walkers))

            # we use a second loop for spliting so we're not bothered by cases
            # like a walker not spliting because we hit the population cap even
            # though the next one would have died and freed a place
            # not a big if it happened though
            for walker in walkers:
                if len(next_walker_pop) < self.params["max_walkers"]:
                    if random() < self.params["split_chance"]:
                        next_walker_pop.append(walker.split())
            walkers = next_walker_pop

        start_x, start_y = -1, -1
        while grid[start_y][start_x] != Tile.FLOOR or start_x == -1:
            start_x, start_y = randint(0, width - 1), randint(0, height - 1)

        result = Map(width, height, grid, start_y, start_x)

        # post-processing: add walls
        for x in range(width):
            for y in range(height):
                if grid[y][x] == Tile.EMPTY:
                    c = sum([1 if grid[j][i] == Tile.FLOOR else 0 for j, i in Map.neighbourhood(grid, y, x, large=True)])
                    if c == 4 and self.params["no_lone_walls"]:
                        result.tiles[y][x] = Tile.FLOOR
                    elif c > 0:
                        result.tiles[y][x] = Tile.WALL
        for x in range(width):
            for y in [0, height-1]:
                if grid[y][x] == Tile.FLOOR:
                    grid[y][x] = Tile.WALL
        for y in range(height):
            for x in [0, width-1]:
                if grid[y][x] == Tile.FLOOR:
                    grid[y][x] = Tile.WALL

        return result