#!/usr/bin/env python
#
# Har logo - always in flux
#
# python-pygame needs to be installed
#

import pygame, sys, random, math

from pygame.locals import *

def G(x, y, sigma):
    return (math.e**-((x*x+y*y)/(2.*sigma*sigma))) / (2.*math.pi*sigma*sigma)

def gauss_blurmatrix( radius, sigma ):
    return [ [ G(x, y, sigma) for x in xrange(1-radius, radius) ]
                              for y in xrange(1-radius, radius) ]

def blur(matrix, blurmatrix):
    h, w = len(matrix), len(matrix[0])
    blur_h, blur_w = len(blurmatrix), len(blurmatrix[0])
    b_m = [ [ 0 for x in xrange(w) ] for y in xrange(h) ]
    xradius = (blur_w-1)/2
    yradius = (blur_h-1)/2
    for y in xrange(h):
        for y2 in xrange(blur_h):
            if 0 <= y+y2-yradius < h:
                for x in xrange(w):
                    for x2 in xrange(blur_w):
                        if 0 <= x+x2-xradius < w:
                            b_m[y+y2-yradius][x+x2-xradius] += \
                                matrix[y][x]*blurmatrix[y2][x2]
    return b_m

def add2d(matrix, a):
    return map(lambda row: map(lambda cell: cell+a, row), matrix)

def get_pos(pos, matrix):
    x, y = pos
    if x < 0 or y < 0 or y >= len(matrix) or x >= len(matrix[y]):
        return 0
    else:
        return matrix[y][x] 

def markov_state_transitions(x, y, matrix):
    total = 0.
    t = []
    v = []
    for d in [ (x-1, y), (x, y-1), (x+1, y), (x, y+1), (x, y),
               (x-1, y-1), (x+1, y-1), (x+1, y+1), (x-1, y+1) ]:
        val = get_pos(d, matrix)
        total += val

        if (val > 0):
            t.append(d)
            v.append(total)

    return map(lambda trans, val: (val/total, trans), t, v)

def markov_chain_transitions(matrix):
    h, w = len(matrix), len(matrix[0])
    return [ [ markov_state_transitions(x, y, matrix)
               for x in xrange(len(matrix[y])) ]
               for y in xrange(len(matrix))    ]
    

def weighed_choice(tuple_list):
    r = random.random()
    for (i, var) in tuple_list:
        if i >= r:
            return var

def next_pos(pos, transitions):
    x, y = pos
    return weighed_choice(transitions[y][x])

def get_square(blocksize):
    #color = [ random.randrange(127, 255) ] * 3
    #c =  random.randrange(32, 127)
    #color = [ c, c, c*2 ]
    color = [ random.randrange(0, 255) for i in xrange(3) ]
    square = pygame.Surface( (blocksize, blocksize) )
    square = square.convert()
    square.fill( color )
    return square

def scale(m, n):
    return [ [ m[y/n][x/n] for x in xrange(len(m[y/n])*n) ]
                           for y in xrange(len(m)*n)    ]

size = 640, 480
title = "HAR2009 logo"
bgcolor = (0, 0, 24)
#n = 400
#gridsize = 8
#blocksize = 9
#scalefactor = 6
#blurradius = 4
#sigma = 1
#n_steps = 4

n = 800
gridsize = 6
blocksize = 8
scalefactor = 8
blurradius = 4
sigma = 1
n_steps = 6

# takes a while to load :-)
#n = 2000
#gridsize = 3
#blocksize = 2
#scalefactor = 16
#blurradius = 8
#sigma = 3
#n_steps = 3

har = (
    ( 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ),
    ( 0, 1, 0, 1, 0, 0, 1, 0, 0, 1, 1, 0, 0 ),
    ( 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0 ),
    ( 0, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 0, 0 ),
    ( 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0 ),
    ( 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0 ),
    ( 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ),
)

def get_har_logo(scalefactor, blurradius, sigma):
    logo = scale(add2d(har, .001), scalefactor)
    return blur(logo, gauss_blurmatrix(blurradius, sigma))

def interpolate(pos1, pos2, frac):
    x1, y1 = pos1
    x2, y2 = pos2
    return ( x2*frac + x1*(1-frac), y2*frac+y1*(1-frac) )

def get_positions(states):
    return map(lambda (x, y): (10+x*gridsize, 60+y*gridsize), states)

def main():

    print "blurring logo"
    logo = get_har_logo(scalefactor, blurradius, sigma)
    print "calculating state transitions"
    t = markov_chain_transitions(logo)

    pygame.init()
    screen = pygame.display.set_mode(size, pygame.FULLSCREEN)
    pygame.display.set_caption(title)

    states = [ (random.randrange(len(logo[0])),
                random.randrange(len(logo)) ) for i in xrange(n) ]
    squares = [ get_square(blocksize) for i in xrange(n) ]

    positions = get_positions(states)

    back = pygame.Surface(screen.get_size())
    back = back.convert()
    back.fill( bgcolor )

    clock = pygame.time.Clock()

    while 1:

        oldpositions = positions[:]
        for i in xrange(n):
            states[i] = next_pos(states[i], t)

        positions = get_positions(states)

        for step in xrange(n_steps):
            clock.tick(60)
            screen.blit(back, (0, 0))

            for i in xrange(n):
                screen.blit(squares[i], interpolate(oldpositions[i], positions[i], float(step)/n_steps))

            pygame.display.flip()

            for event in pygame.event.get():
                if event.type in (QUIT, KEYDOWN):
                    return

if __name__ == '__main__':
    main()