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

DIRS = ((0, +1),
        (0, -1),
        (-1, 0),
        (+1, 0))

def create_maze(mw, mh):
    """
    assume maze is a map,
    use depth first search to connect the map,
    and so is a maze.
    """
    path = [(0,0)] #the search path
    vector_link = [] #links between vectors, for draw maze
    visited = np.zeros((mw,mh)) #for O(1) search
    visited[0,0] = 1

    while True:
        #get next vector
        nextp = get_next(path, visited, mw, mh)
        #check if this vector is the end
        if not nextp:
            # 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[nextp] = 1
        path.append(nextp)

    #create maze
    maze = np.ones((mw*2-1, mh*2-1), dtype=np.int)
    for start, end in vector_link:
        #print start, end
        d1 = start[0]*2, start[1]*2
        d2 = end[0]*2, end[1]*2
        middle = ((start[0] + end[0]),
                  (start[1] + end[1]))
        maze[d1] = 0
        maze[d2] = 0
        maze[middle] = 0
    return maze

def get_next(path, visited, mw, mh):
    """
    get what is the next vector?
    """
    head = path[-1]
    nextps = [] #store next vectors
    for mod in DIRS:
        newp = head[0] + mod[0], head[1] + mod[1]
        #out of the scope
        if ((0 > newp[0]) or (mw <= newp[0]) or
            (0 > newp[1]) or (mh <= newp[1])
            ):
            continue
        #have visited
        if visited[newp] == 1:
            continue
        nextps.append(newp)
    if nextps == []:
        return None
    #random choose a vector
    return nextps[random.randint(0, len(nextps)-1)]

def show_maze(maze):
    """
    print maze to command line.
    """
    for row in maze:
        print "".join(["%d"%i for i in row])

def draw_maze(maze):
    """
    use pygame to render a png file.
    """
    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] == 0:
                surface.fill((200,200,200), 
                             pygame.Rect(i*DOT, j*DOT, DOT, DOT))
    pygame.image.save(surface, "test.png")

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

if __name__=="__main__":
    main()