updated spinner.py script

from collections import defaultdict
from pprint import pprint

# From http://code.activestate.com/recipes/578231-probably-the-fastest-memoization-decorator-in-the-/
def memodict(f):
    """ Memoization decorator for a function taking a single argument """
    class memodict(dict):
        def __missing__(self, key):
            ret = self[key] = f(key)
            return ret
    return memodict().__getitem__

# We assume that stuck spinners will never move again
# We mark this case with tick count = 0
# 1 tick = 2 frames

# Helper function to compute base distribution when spinner tick count reaches 0
def tick_distribution():
    out = {}
    for x in range(32):
        out[x] = 1.0 / 32
    return out
base_tick_distribution = tick_distribution()

DOWN = 0
UP = 1
LEFT = 2
RIGHT = 3
# Helper function to compute probability of new direction given old direction
def direction_distribution():
    out = {}
    for direction in range(4):
        t = {}
        opposite = (direction + 2) % 4
        t[3 - direction] = 1.0 / 2
        t[opposite] = 1.0 / 4
        t[3 - opposite] = 1.0 / 4
        out[direction] = t
    return out
base_direction_distribution = direction_distribution()

# Given current state, compute distribution for next state
# State is of the form (direction, n_ticks)
@memodict
def next_state(state):
    out = {}
    if state[1] != 1:
        # Decrement n_ticks
        # Handle n_ticks == 0 case to not change state
        new_state = (state[0], max(0, state[1] - 1))
        out[new_state] = 1
    else:
        # Use the base distributions computed earlier to get new distribution
        new_directions = base_direction_distribution[state[0]]
        new_ticks = base_tick_distribution
        for direction, p in new_directions.iteritems():
            for n_ticks, q in base_tick_distribution.iteritems():
                new_state = (direction, n_ticks)
                out[new_state] = p * q
    return out

# Given current state distribution, compute next state distribution
def next_distribution(distribution):
    out = defaultdict(float)
    for state, p in distribution.iteritems():
        d = next_state(state)
        for new_state, q in d.iteritems():
                out[new_state] += p * q
    return out

# Computes distribution given observation of spinner in given direction for n ticks
# Possible values of initial hidden state n_ticks: n_ticks > n or n_ticks == 0
# If n >= 31, then the spinner is guaranteed to be stuck.
def start_distribution(direction, n):
    out = {}
    n_possibilities = max(1, 32 - n)
    for i in range(n_possibilities):
        out[(direction, i)] = 1.0 / n_possibilities
    return out

# Remove hidden n_ticks part of state
# Compute distribution based only on direction
def compress(dist):
    out = [0.0, 0.0, 0.0, 0.0]
    for state, p in dist.iteritems():
            out[state[0]] += p
    return out

# Given start state corresponding to d, n
# After up to t ticks, computes direction distribution
def chain(d, n, t):
    out = []
    s = start_distribution(d, n)
    for i in range(t+1):
        out.append(compress(s))
        s = next_distribution(s)
    return out

def main():
    # Collect some probabilities and dump them in files
    # Yes this is ugly but i dont care
    bike_file = open('bike.txt', 'w')
    walk_file = open('walk.txt', 'w')
    bike_file.write('\t| down\t| up\t| left\t| right\n')
    bike_file.write('------------------------------------------\n')
    walk_file.write('\t| down\t| up\t| left\t| right\n')
    walk_file.write('------------------------------------------\n')
    # Without loss of generality, suppose that we cross below the spinner
    for n_ticks in range(32):
        # Multiply by 2 to convert to frames
        bike_file.write('%d' % (2*n_ticks))
        walk_file.write('%d' % (2*n_ticks))
        for start_direction in range(4):
            bike_file.write('\t| ')
            walk_file.write('\t| ')
            x = chain(start_direction, n_ticks, 9)
            # x[k][DOWN] = probability of spinner facing down after k ticks
            # 8 frames = 4 ticks per bike, yet for some reason it actually takes 5...
            bike_file.write('%.3f' % x[5][DOWN])
            # 16 frames = 8 ticks per step, yet for some reason it actually takes 9...
            walk_file.write('%.3f' % x[9][DOWN])
        # so much copypasta
        bike_file.write('\n')
        walk_file.write('\n')
    bike_file.close()
    walk_file.close()

if __name__ == '__main__':
    main()