#!/usr/bin/python3
"""Quick maze generator program for 2-D grids.

Specifically I wanted to build a 2-D maze in Minetest, so I hacked
together this program in about 20 minutes, using maze-generation
algorithms I already knew.  It outputs a random ASCII-art maze with a
single valid path that looks like this:

      0  1  2  3  4  5  6  7  8  9 10 11 12 13 14
  0  @@    @@ @@ @@ @@ @@ @@ @@ @@ @@ @@ @@ @@ @@
  1  @@    @@    @@          @@                @@
  2  @@    @@    @@ @@ @@    @@    @@ @@ @@ @@ @@
  3  @@    @@          @@    @@          @@    @@
  4  @@    @@    @@ @@ @@    @@    @@    @@    @@
  5  @@                      @@    @@          @@
  6  @@ @@ @@ @@ @@ @@ @@    @@    @@ @@ @@ @@ @@
  7  @@                                        @@
  8  @@ @@ @@ @@ @@    @@ @@ @@ @@ @@    @@    @@
  9  @@                            @@    @@    @@
 10  @@    @@    @@    @@    @@ @@ @@ @@ @@    @@
 11  @@    @@    @@    @@          @@    @@    @@
 12  @@ @@ @@ @@ @@    @@ @@ @@    @@    @@ @@ @@
 13  @@                      @@                @@
 14  @@ @@ @@ @@ @@ @@ @@ @@ @@ @@ @@ @@ @@    @@

On my laptop it takes about 9.4 seconds to generate a 1701×1701 maze,
which is 8.3 megabytes; 3401×1701 is 21.0 seconds, and 3401×3401 is
50.7 seconds (and about 1.8 gigabytes of RAM), so it’s measurably
worse than linear time, but not ridiculously so.  I haven’t checked
that its results are correct for such large mazes, and definitely at
least the row and column numbers will be cattywampus, but for smaller
ones they look okay.

The objective here is to make it convenient to reproduce the maze in
some other medium, so row and column numbers are provided for
reference.

"""
import sys, random


def make_maze(size):
    width, height = size
    assert width % 2 == 1       # 3, 5, 7, etc., but not 2, 4, etc.
    assert height % 2 == 1      # 3, 5, 7, etc., but not 2, 4, etc.
    full = [[1 for _ in range(width)] for _ in range(height)]
    for row in range(1, height, 2):
        for col in range(1, width, 2):
            full[row][col] = 0

    # Entry and exit.
    full[0][1] = 0
    full[-1][-2] = 0

    # Enumerate internal walls so we can shuffle the order to remove them in.
    walls = [(wr, wc, pyramis, thisbe)
             for row in range(1, height, 2)
             for col in range(1, width, 2)
             for wr, wc, pyramis, thisbe in [(row-1, col, (row-2, col), (row, col)),
                                             (row, col-1, (row, col-2), (row, col))]
             if wr and wc]

    random.shuffle(walls)

    # Union-find algorithm so we can remove only walls that don’t
    # create loops.
    uf = {}
    def root_of(cell):
        if cell not in uf:
            return cell
        uf[cell] = result = root_of(uf[cell])
        return result

    for wr, wc, pyramis, thisbe in walls:
        pr, tr = root_of(pyramis), root_of(thisbe)
        if pr != tr:
            # The cells were previously disconnected; erase the wall
            # and mark them as connected.
            uf[min(pr, tr)] = max(pr, tr)
            full[wr][wc] = 0

    return full

def print_maze(cells):
    print(' ' * 4, end='')
    for col, _ in enumerate(cells[0]):
        print('{:3d}'.format(col), end='')
    print()

    for rownum, row in enumerate(cells):
        print('{:3d} '.format(rownum), end='')
        for cell in row:
            print(' @@' if cell else '   ', end='')

        print()

def main(args):
    print_maze(make_maze((int(args[1] if len(args) > 1 else 15),
                          int(args[2] if len(args) > 2 else 15))))

if __name__ == '__main__':
    main(sys.argv)