#!/usr/bin/python
'''
Jon Beckham <beckham@leftorium.net>, code sample

Generates a minimally spanning maze, thickening the walls to match an input
image.

Samples available at http://strok.in/mazer

Requires the Python Imaging Library.
  http://www.pythonware.com/library/pil/handbook/index.htm
'''

import optparse
import random
import sys

import Image
import ImageDraw

X, Y = range(2)
TOP, RIGHT, BOTTOM, LEFT = range(4)

class Cell(object):
    __slots__ = ('pos', 'size', 'walls')

    def __init__(self, pos, size):
        self.pos = (pos[X] * size, pos[Y] * size)
        self.size = size
        self.walls = []

    def setwalls(self, walls):
        self.walls = walls
        self.walls[TOP].set_endpoints( [self.tl(), self.tr()] )
        self.walls[RIGHT].set_endpoints( [self.tr(), self.br()] )
        self.walls[BOTTOM].set_endpoints( [self.bl(), self.br()] )
        self.walls[LEFT].set_endpoints( [self.tl(), self.bl()] )

    def tl(self):
        return self.pos

    def tr(self):
        return (self.pos[X] + self.size, self.pos[Y])

    def bl(self):
        return (self.pos[X], self.pos[Y] + self.size)

    def br(self):
        return (self.pos[X] + self.size, self.pos[Y] + self.size)

    def draw_walls(self, drawer, data):
        [wall.draw(drawer, data) for wall in self.walls]


class Wall(object):
    __slots__ = ('cells', 'endpoints', 'dir', 'perp', 'len', 'up', 'drawn', 'probe')

    def __init__(self, cells):
        if 2 != len(cells):
            raise ValueError('a wall can only have two cells.')
        self.cells = cells
        self.endpoints = ()
        self.dir = None
        self.perp = None
        self.len = 0
        self.up = True
        self.drawn = False

    def __nonzero__(self):
        return self.up

    def set_endpoints(self, pointlist):
        # ensure points are ordered to reduce complexity later
        if self.endpoints:
            # when a second cell neighboring this wall is initialized, this is
            # called again.  noop in that case.
            return
        pointlist.sort()
        self.endpoints = pointlist
        for d in (X, Y):
            self.len += (self.endpoints[1][d] - self.endpoints[0][d])
            if self.endpoints[0][d] == self.endpoints[1][d]:
                self.perp = d
            else:
                self.dir = d
        self.probe = (self.len - 1) / 2  # probe * 2 < len

    def remove(self):
        if not self:
            raise ValueError('cannot re-remove a removed wall.')
        if None not in self.cells:
            # not an edge wall, so we can remove it. (noop on edge walls.)
            self.up = False

    def addcell(self, newcell):
        if newcell in self.cells:
            raise ValueError("%s already in my neighboring cell list...")
        if not self.cells[0]:
            self.cells[0] = newcell
        elif not self.cells[1]:
            self.cells[1] = newcell

    def on_edge(self):
        return None in self.cells

    def draw(self, drawer, data):
        if self and not self.drawn:
            self.weighted(drawer, data)
            self.drawn = True

    ### Drawing ###
    def _weight_walk(self, data, orientation, pos, start, direction, limit):
        'walk from start to limit, returning the sum of the pixels traversed'
        pixelsum = 0

        if direction < 0:
            edgefunc = max
        else:
            edgefunc = min
        end = edgefunc(limit, start + direction*self.probe)
        for i in range(start, end, direction):
            if orientation == X:
                inpixel = data[pos][i]
            elif orientation == Y:
                inpixel = data[i][pos]
            pixelsum += 255 - inpixel
        return pixelsum

    def weighted(self, drawer, data):
        'Draw the maze with solid walls thickened based on input data.'
        maxedge = (len(data[0]), len(data))

        # extend the wall's length to represent the input pixels
        psum = self._weight_walk(data, self.dir, self.endpoints[0][self.perp], self.endpoints[0][self.dir], -1, 0)
        start = self.endpoints[0][self.dir] - psum / 255
        psum = self._weight_walk(data, self.dir, self.endpoints[1][self.perp], self.endpoints[1][self.dir], 1, maxedge[self.dir])
        end = self.endpoints[1][self.dir] + psum / 255

        # thicken the wall to represent the input pixels
        perp_pos = self.endpoints[0][self.perp]
        for i in range(start, end):
            psum = self._weight_walk(data, self.perp, i, perp_pos, -1, 0)
            pstart = perp_pos - psum / 255
            psum = self._weight_walk(data, self.perp, i, perp_pos, 1, maxedge[self.perp])
            pend = perp_pos + psum / 255

            if self.dir == X:
                drawer.line(((i, pstart), (i, pend+1)), fill=0)
            elif self.dir == Y:
                drawer.line(((pstart, i), (pend, i)), fill=0)

        # even if there wasn't any weight in the input, we still want a wall
        drawer.line(self.endpoints, fill=0)


class Maze(object):
    def __init__(self, gridsize, cellsize, inputfile, outputfile):
        self.gridsize = gridsize  # (width, height)
        self.cellsize = cellsize
        self.inputfile = inputfile
        self.outputfile = outputfile

        self.grid = []
        for row in range(self.gridsize[Y]):
            self.grid += [[]]
            for col in range(self.gridsize[X]):
                self.grid[row] += [[]]
                newcell = Cell((col, row), self.cellsize)

                if 0 <= row - 1:
                    w0 = self.grid[row-1][col].walls[BOTTOM]
                    w0.addcell(newcell)
                else:
                    w0 = Wall([None, newcell])

                w1 = Wall([newcell, None])

                w2 = Wall([newcell, None])

                if 0 <= col - 1:
                    w3 = self.grid[row][col-1].walls[RIGHT]
                    w3.addcell(newcell)
                else:
                    w3 = Wall([None, newcell])

                newcell.setwalls([w0, w1, w2, w3])
                self.grid[row][col] = newcell

        self.imgsize = (self.gridsize[X] * self.cellsize + 1, self.gridsize[Y] * self.cellsize + 1)

    def generate(self):
        '''Use Prim's algorithm to generate a minimally spanning maze.'''
        c = random.choice(random.choice(self.grid))
        maze = set([c,])
        walls = [w for w in c.walls if w and not w.on_edge()]

        while walls:
            random_wall = random.choice(walls)
            for neighbor in random_wall.cells:
                if neighbor not in maze:
                    maze.add(neighbor)
                    random_wall.remove()
                    walls += [w for w in neighbor.walls if w and not w.on_edge()]
            walls.remove(random_wall)
            walls = filter(None, walls)  # remove all walls from consideration that aren't up

    def draw(self):
        outimage = Image.new(mode='L', size=self.imgsize, color=255)
        drawer = ImageDraw.Draw(outimage)

        indata = []
        if self.inputfile:
            inimage = Image.open(self.inputfile)
            inimage = inimage.convert('L')

            # crop/resize the input into the size and aspect of the output
            outaspect = 1.0 * outimage.size[X] / outimage.size[Y]
            inaspect = 1.0 * inimage.size[X] / inimage.size[Y]
            xratio = 1.0 * inimage.size[X] / outimage.size[X]
            yratio = 1.0 * inimage.size[Y] / outimage.size[Y]
            if inaspect > outaspect:
                offset = abs(inimage.size[X] - outimage.size[X] * yratio) / 2
                extents = (offset, 0, inimage.size[X] - offset, inimage.size[Y])
                inimage = inimage.transform(outimage.size, Image.EXTENT, extents)
            elif inaspect < outaspect:
                offset = abs(inimage.size[Y] - outimage.size[Y] * xratio) / 2
                extents = (0, offset, inimage.size[X], inimage.size[Y] - offset)
                inimage = inimage.transform(outimage.size, Image.EXTENT, extents)
            else:
                inimage = inimage.resize(outimage.size)

            # getdata() gives a flat "list"; make it into a 2-D list()
            # (this is still faster than getpixel())
            flatdata = list(inimage.getdata())
            indata = [flatdata[i:i+self.imgsize[X]] for i in range(0, len(flatdata), self.imgsize[X])]

        for row in self.grid:
            for cell in row:
                cell.draw_walls(drawer, indata)

        del drawer
        outimage.save(self.outputfile, 'PNG')


def main():
    parser = optparse.OptionParser()
    parser.set_usage('%prog [options]')
    parser.add_option('-s', '--size', action='store', default='30x30',
        help='set grid size of resultant maze [default: 30x30]')
    parser.add_option('-c', '--cellsize', action='store', type='int', default='10',
        help='set size of maze cells [default: 10]')
    parser.add_option('-i', '--input', action='store',
        help='image used for wall weighting')
    parser.add_option('-o', '--output', action='store', default='output.png',
        help='resultant image file')
    (opts, args) = parser.parse_args()

    fail = False
    if not 'x' in opts.size:
        print >>sys.stderr, 'Invalid argument to --size; expected <int>x<int>.'
        fail = True
    size = [int(n) for n in opts.size.split('x')]

    if not opts.input:
        print >>sys.stderr, '--input required.'
        fail = True

    if fail:
        parser.print_help()
        return 1

    m = Maze(size, opts.cellsize, opts.input, opts.output)
    m.generate()
    m.draw()

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