root/confluence/foreverbubble/foreverbubble.py

#!/usr/bin/python
#
# requires: argparse, networkx.  Works on Linux.  Might work on OSX.  Probably doesn't work on Windows.  Patches appreciated.
#
# Copyright 2010 Adrianna Pinska, Simon Cross
# This program is distributed under the terms of the GNU General Public License v3

import warnings
warnings.simplefilter("ignore", DeprecationWarning)

import random
import os
from copy import deepcopy

from argparse import ArgumentParser
from networkx import Graph, connected_components

class Ball(object):
    def __init__(self, ball=None):
        self.ball = ball

    def __str__(self):
        if self.ball is None:
            return '-'
        else:
            return str(self.ball)


class Row(object):
    def __init__(self, row_index):
        self.even = not row_index % 2
        self.balls = []

    def num_balls(self):
        return 8 if self.even else 7

    def __str__(self):
        prefix = '  ' if not self.even else ''
        return prefix + '   '.join([str(b) for b in self.balls])


class RandomRow(Row):
    def __init__(self, row_index, selection, symmetric):
        super(RandomRow, self).__init__(row_index)
        if symmetric:
            startballs = [Ball(random.choice(selection)) for b in range(4)]
            self.balls = startballs + list(reversed(deepcopy(startballs)[:self.num_balls()/2]))
        else:
            self.balls = [Ball(random.choice(selection)) for b in range(self.num_balls())]

class BlankRow(Row):
    def __init__(self, row_index):
        super(BlankRow, self).__init__(row_index)
        self.balls = [Ball()] * self.num_balls()

class Level(object):
    BALLS = range(8)

    ASYM, SYM_HORIZ, SYM_BOTH = range(3)
    SYMMETRIES = {
        "asym" : ASYM,
        "horiz" : SYM_HORIZ,
        "both" : SYM_BOTH,
    }

    def __init__(self, fill_rows=None, colours=None, emptyspace=None, symmetry=None):
        if fill_rows is None:
            fill_rows = random.randint(6,10)
        if colours is None:
            colours = random.randint(1,8)
        if emptyspace is None:
            emptyspace = random.randint(1,4)
        if symmetry is None:
            symmetry = random.randint(0,2)

        self.rows = []

        selection = random.sample(self.BALLS, colours) + [None] * emptyspace
        symmetric_rows = symmetry > 0

        if symmetry > 1:
            if not fill_rows % 2:
                fill_rows -= 1
            for row in range((fill_rows + 1)/2):
                self.rows.append(RandomRow(row, selection, symmetric_rows))
            for rand_row in reversed(self.rows[:fill_rows/2]):
                self.rows.append(deepcopy(rand_row))
        else:
            for row in range(fill_rows):
                self.rows.append(RandomRow(row, selection, symmetric_rows))

        for row in range(fill_rows, 10):
            self.rows.append(BlankRow(row))

    def valid(self):
        g = Graph()

        def add_edge(b1, b2):
            if b1.ball is not None and b2.ball is not None:
                g.add_edge(b1, b2)

        # add all balls and horizontal edges
        for row in self.rows:
            for ball in row.balls:
                if ball.ball is not None:
                    g.add_node(ball)
            for b1, b2 in zip(row.balls[:-1], row.balls[1:]):
                add_edge(b1, b2)

        # add inter-row connections
        for row, above, below in zip(self.rows[1::2], self.rows[::2], self.rows[2::2] + [BlankRow(0)]):
            for ball, above_left, above_right, below_left, below_right in zip(row.balls, above.balls[:-1], above.balls[1:], below.balls[:-1], below.balls[1:]):
                add_edge(ball, above_left)
                add_edge(ball, above_right)
                add_edge(ball, below_left)
                add_edge(ball, below_right)

        top_row = set(self.rows[0].balls)
        valid = True
        for comp in connected_components(g):
            if set(comp).isdisjoint(top_row):
                valid = False
                break

        return valid

    def __str__(self):
        return '\n'.join([str(r) for r in self.rows])


class LevelSet(object):
    def __init__(self, num_levels, **kwargs):
        l = 0
        self.levels = []

        while l < num_levels:
            level = Level(**kwargs)
            if level.valid():
                self.levels.append(level)
                l += 1

    def __str__(self):
        return '\n\n'.join([str(l) for l in self.levels])

if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('-f', '--filename', dest='filename', type=str, default="%s/.frozen-bubble/levels/foreverbubble" % os.getenv('HOME'), metavar='FILE', help='levelset save location (default: %(default)s)')
    parser.add_argument('-n', '--number', dest='number', type=int, default=20, metavar='NUMBER', help='number of levels to generate (default: %(default)s)')

    parser.add_argument('-r', '--rows', dest='fill_rows', type=int, choices=xrange(1,11), metavar='ROWS', help='number of rows to generate per level (permitted values: %(choices)s; default: random)')
    parser.add_argument('-c', '--colours', dest='colours', type=int, choices=xrange(1,9), metavar='COLOURS', help='number of colours to use per level (permitted values: %(choices)s; default: random)')
    parser.add_argument('-e', '--empty-space', dest='emptyspace', type=int, choices=xrange(0,4), metavar='SPACE', help='amount of empty space on each level (permitted values: %(choices)s; default: random)')
    parser.add_argument('-s', '--symmetry', dest='symmetry', type=str, choices=Level.SYMMETRIES.keys(), metavar='SYMMETRY', help='symmetry of each level (permitted values: %(choices)s; default: random)')
    args = parser.parse_args()

    symmetry = Level.SYMMETRIES[args.symmetry] if args.symmetry else None

    levelset = LevelSet(args.number, fill_rows=args.fill_rows, colours=args.colours, emptyspace=args.emptyspace, symmetry=symmetry)

    levelsetfile = open(args.filename, "w")
    levelsetfile.write(str(levelset))
    levelsetfile.close()