Opus Magnum Information Age puzzle generator (Python 2)

opus_magnum.py

import sys
from struct import pack, unpack
import random


class Primes:
    Salt         = 0x1
    Air          = 0x2
    Earth        = 0x3
    Fire         = 0x4
    Water        = 0x5
    Mercury      = 0x6
    Gold         = 0x7
    Silver       = 0x8
    Copper       = 0x9
    Iron         = 0xa
    Tin          = 0xb
    Lead         = 0xc
    Life         = 0xd
    Death        = 0xe
    Repeat       = 0xf
    Quintessence = 0x10

BOND = 0x1
TRIPLEX = 0xe


def decode_varint(f):
    size = unpack("B", f.read(1))[0]
    if size >= 128:
        raise NotImplementedError
    return size

def decode_string(f):
    length = decode_varint(f)
    return f.read(length)

def decode_int(f):
    return unpack("i", f.read(4))[0]

def decode_molecule(f):
    print "-"*79
    for i in xrange(decode_int(f)):   # primes
        el, x, y = unpack("bbb", f.read(3))
        for name, value in Primes.__dict__.items():
            if el == value:
                el = name
        print el, x, y
    for i in xrange(decode_int(f)):   # bonds
        bt, sx, sy, dx, dy = unpack("bbbbb", f.read(5))
        bonds = []
        if bt & 1: bonds.append('bond')
        if bt & 2: bonds.append('red')
        if bt & 4: bonds.append('black')
        if bt & 8: bonds.append('yellow')
        print '+'.join(bonds), sx, sy, dx, dy
    print "-"*79

def decode(filename):
    f = open(filename, 'rb')
    version = decode_int(f)
    assert version == 3, version
    print repr(decode_string(f))
    creator = f.read(8)
    tools = unpack("q", f.read(8))[0]
    print hex(tools)
    for i in xrange(decode_int(f)):   # inputs
        decode_molecule(f)
    for i in xrange(decode_int(f)):   # outputs
        decode_molecule(f)
    print 'target scale:', decode_int(f)
    print 'is production:', ord(f.read(1))


def encode_int(x):
    return pack("i", x)

def encode_string(x):
    assert len(x) < 128
    return chr(len(x)) + x

def encode_atom(el, x, y):
    return pack("bbb", el, x, y)

def encode_bond(bond, sx, sy, dx, dy):
    return pack("bbbbb", bond, sx, sy, dx, dy)

class Molecule:
    def __init__(self):
        self.atoms = []
        self.bonds = []

    def produce(self, g):
        g.write(encode_int(len(self.atoms)))
        for atom in self.atoms:
            g.write(encode_atom(*atom))
        g.write(encode_int(len(self.bonds)))
        for bond in self.bonds:
            g.write(encode_bond(*bond))

def produce(g):
    g.write(encode_int(3))
    g.write(encode_string("INFORMATION AGE"))
    g.write('!' * 8)
    g.write(encode_int(0x7c2470f) + encode_int(0))
    #
    # inputs
    numbers = [random.randrange(0, 6) for _ in range(42)]
    g.write(encode_int(3))

    m = Molecule()
    for x in range(len(numbers)):
        if x > 0:
            m.bonds.append((BOND, x - 1, 0, x, 0))
        m.atoms.append((Primes.Salt, x, 0))
        for j in range(numbers[len(numbers)-1-x]):
            m.atoms.append((Primes.Air, x, j + 1))
            m.bonds.append((BOND, x, j, x, j + 1))
    m.produce(g)

    m = Molecule()
    m.atoms.append((Primes.Lead, 0, 0))
    m.produce(g)

    m = Molecule()
    m.atoms.append((Primes.Mercury, 0, 0))
    m.produce(g)

    #
    # outputs
    g.write(encode_int(1))
    m = Molecule()
    for x in range(len(numbers)):
        if x > 0:
            m.bonds.append((BOND, -x+1, 0, -x, 0))
        m.atoms.append((Primes.Lead - numbers[x], -x, 0))
    m.produce(g)
    #
    g.write(encode_int(1))
    g.write('\x00')


if __name__ == '__main__':
    if len(sys.argv) > 1:
        decode(sys.argv[1])
    else:
        assert not sys.stdout.isatty(), "redirect the output to a file"
        produce(sys.stdout)