#!/usr/bin/env python
#-*- coding:utf-8 -*-
"""
use map depth first search to create a maze.
this map is multi dimension!
"""
import random
import numpy as np

def create_maze(*size):
    """
    assume maze is a map,
    use depth first search to connect the map,
    and so is a maze.
    """
    dim = len(size) #dimension
    start = np.zeros(dim, dtype=np.int)
    path = [start] #the search path
    vector_link = [] #links between vectors, for draw maze
    visited = np.zeros(size, dtype=np.bool) #for O(1) search
    visited[tuple(start.tolist())] = True
    #directions
    DIRS = []
    for i in range(dim):
        z = np.zeros(dim, dtype=np.int)
        z[i] += 1
        DIRS.append(z)
        z = np.zeros(dim, dtype=np.int)
        z[i] -= 1
        DIRS.append(z)

    while True:
        #get next vector
        nextp = get_next(path, visited, size, DIRS)
        #check if this vector is the end
        if nextp is None:
            # if so, pop to previous head
            path.pop()
            # if path is empty, search is finished
            if path == []:
                #finish search
                break
            continue
        #store data, move to next head
        vector_link.append((path[-1],nextp))
        visited[tuple(nextp.tolist())] = 1
        path.append(nextp)

    #create maze
    maze_size = np.array(size)*2 -1
    maze = np.zeros(maze_size, dtype=np.bool)
    for start, end in vector_link:
        #print start, end
        d1 = start*2
        d2 = end*2
        middle = start + end
        maze[tuple(d1.tolist())] = True
        maze[tuple(d2.tolist())] = True
        maze[tuple(middle.tolist())] = True
    return maze

def get_next(path, visited, size, DIRS):
    """
    get what is the next vector?
    """
    head = path[-1]
    #print "head:",head
    nextps = [] #store next vectors
    for mod in DIRS:
        newp = head + mod
        #print 'newp:',newp
        #out of the scope
        out_scope = False
        for i, v in enumerate(newp):
            if (v<0) or (v>=size[i]): 
                #print newp, v, size[i], size
                out_scope = True
                break
        if out_scope:
            continue
        #have visited
        if visited[tuple(newp.tolist())] == True:
            # print "visited:",newp
            # print visited
            continue
        nextps.append(newp)
    if nextps == []:
        return None
    #random choose a vector
    nextp = nextps[random.randint(0, len(nextps)-1)]
    #print "nextps:",nextps
    #print "nextp:  ",nextp
    return nextp

def show_maze(maze):
    """
    print maze to command line.
    """
    if maze.ndim == 3:
        for map in maze:
            print
            for row in map:
                print "".join(['_'if i else '#' for i in row])
    elif maze.ndim == 2:
        for row in maze:
            print "".join(['_'if i else '#' for i in row])

def draw_maze(maze):
    """
    use pygame to render a png file.
    """
    if maze.ndim != 2: return
    try:
        import pygame
    except:
        print "looks like you don't have pygame,"
        print "install pygame for rending a image file."
    pygame.init()

    DOT = 8
    size = maze.shape
    sw, sh = size

    surface = pygame.Surface((sw*DOT, sh*DOT))
    for i in range(sw):
        for j in range(sh):
            if maze[i][j] == True:
                surface.fill((200,200,200), 
                             pygame.Rect(i*DOT, j*DOT, DOT, DOT))
    pygame.image.save(surface, "test.png")

def main():
    maze = create_maze(30, 30, 30)
    show_maze(maze)
    draw_maze(maze)

if __name__=="__main__":
    main()