#!/usr/bin/python3
#
# Messing with maze generation. This uses:
# https://en.wikipedia.org/wiki/Maze_generation_algorithm#Wilson's_algorithm
#

import sys
import random
import pprint

# Each cell has four walls, organized NESW (index: 0..3) as True/False
# Maze is a DICT: COORD: [N,E,S,W]
# 0,0 is upper-left (North/West) on a printed maze


def gen_maze(width, height):
    #print(_which_dir((1,2),(1,3)))
    #print(_which_dir((1,2),(0,2)))
    #print(_which_dir((1,2),(1,1)))
    #print(_which_dir((1,2),(2,2)))
    #return

    maze = { }
    notmaze = set((x, y) for x in range(width) for y in range(height))

    # Choose an initial location, make it part of the maze
    initial = _pick(notmaze)
    maze[initial] = [True]*4
    notmaze.remove(initial)
    print('INITIAL:', initial)

    while notmaze:
        #print_maze(maze, width, height)

        # Store coordinates ordered along a random walk.
        walk = [ ]

        # Generate a random walk.
        while True:
            if len(walk) == 0:
                # Pick a starting point.
                walk.append(_pick(notmaze))

            current = walk[-1]
            #print('CURRENT:', current)
                
            # Find a legal neighbor.
            while True:
                d = random.randrange(4)
                if d == 0:
                    if current[1] == 0:
                        continue
                    neighbor = (current[0], current[1]-1)
                elif d == 1:
                    if current[0] == width-1:
                        continue
                    neighbor = (current[0]+1, current[1])
                elif d == 2:
                    if current[1] == height-1:
                        continue
                    neighbor = (current[0], current[1]+1)
                elif d == 3:
                    if current[0] == 0:
                        continue
                    neighbor = (current[0]-1, current[1])
                break

            # Is this neighbor already in the walk? If so, then remove
            # the whole loop and retry a walk that doesn't loop.
            for i in range(len(walk)):
                if neighbor == walk[i]:
                    #print('REMOVING:', walk[i:])
                    del walk[i:]
                    break
            else:
                # We didn't see a loop. What is at this coordinate?

                if neighbor in maze:
                    # The neighbor is part of the maze, so add the walk.
                    #print('ADDING:', walk, neighbor)
                    current = walk[0]
                    maze[current] = [True]*4
                    notmaze.remove(current)
                    for w in walk[1:]:
                        d = _which_dir(current, w)
                        maze[current][d] = False
                        maze[w] = [True]*4
                        maze[w][d^2] = False
                        notmaze.remove(w)
                        current = w
                    d = _which_dir(current, neighbor)
                    maze[current][d] = False
                    maze[neighbor][d^2] = False

                    # break, to do another walk
                    break
                else:
                    # Just another coordinate. Add to the walk.
                    walk.append(neighbor)

    return maze


def _pick(coords):
    return random.choice(tuple(coords))


def _which_dir(c1, c2):
    #print(f'COORDS: ({c1[0]},{c1[1]}) and ({c2[0]},{c2[1]})')
    if c1[0] == c2[0]:
        return c2[1]-c1[1]+1
    return c1[0]-c2[0]+2


def print_maze(maze, width, height):
    for y in range(height):
        N = [ ]
        EW = [ ]
        for x in range(width):
            c = maze.get((x, y), [True]*4)
            N.append('--' if c[0] else '  ')
            EW.append('|' if c[3] else ' ')
            EW.append('  ')
        print(f'+{"+".join(N)}+')
        print(f'{"".join(EW)}|')
    print(f'+{"--+"*width}')


if __name__ == '__main__':
    width = 10
    height = 10
    m = gen_maze(width, height)
    print_maze(m, width, height)