
import pygame.image
import pygame.draw
import pygame
import math
import collections
from math import pi, sin, cos

vec2 = collections.namedtuple('vec2', 'x y')
vec3 = collections.namedtuple('vec3', 'x y z')

def init(width, height):
    global WIDTH, HEIGHT
    print "Python init function executed."
    WIDTH = width
    HEIGHT = height

BALL_SIZE = 24
ball = pygame.Surface((BALL_SIZE, BALL_SIZE))
ball.fill((0,0,0))
ball.set_colorkey((0,0,0))
for radius in reversed(xrange(BALL_SIZE/2)):
    brightness = (BALL_SIZE/2-radius)/(BALL_SIZE/2.)
    pygame.draw.circle(ball, (int(128*brightness), 
                              int(255*brightness),
                              0), 
                       (BALL_SIZE/2, BALL_SIZE/2),
                       radius, 0)

def generate_box(n_balls, radius):
    balls_per_side = n_balls / 12  
    coords = []
    edge = [-radius + i * (radius * 2.0 / balls_per_side) 
             for i in xrange(balls_per_side)]
    r = radius
    for i in xrange(balls_per_side):
        coords.append((edge[i], -r, -r))
        coords.append((edge[i], r, -r))
        coords.append((edge[i], -r, r))
        coords.append((edge[i], r, r))

        coords.append((-r, edge[i], -r))
        coords.append((r, edge[i], -r))
        coords.append((-r, edge[i], r))
        coords.append((r, edge[i], r))

        coords.append((-r, -r, edge[i]))
        coords.append((r, -r, edge[i]))
        coords.append((-r, r, edge[i]))
        coords.append((r, r, edge[i]))
    coords.append((r,r,r)) # The only corner not yet covered
    return coords

def generate_helix(n_balls, radius, halflength):
    half_balls = n_balls / 2
    return [(radius * sin(4 * pi * i / half_balls),
             radius * cos(4 * pi * i / half_balls),
             -halflength + 2 * halflength * i / half_balls)
            for i in xrange(half_balls)] + \
            [(radius * sin(4 * pi * i / half_balls + 0.6 * pi),
             radius * cos(4 * pi * i / half_balls + 0.6 * pi),
             -halflength + 2 * halflength * i / half_balls)
            for i in xrange(half_balls)] 

def generate_ball(n_balls, n_slices, radius):
    coords = []
    for j in xrange(n_slices):
        z = -radius + j*(2*radius/(n_slices-1))
        slice_radius = math.sqrt(radius ** 2 - z ** 2)
        balls_in_slice = int(1.6*n_balls/n_slices*slice_radius/radius)
        if balls_in_slice == 0:
            balls_in_slice = 1
        for i in xrange(balls_in_slice):
            theta = i * 2 * pi / balls_in_slice
            if z < 0 and n_slices % 2 != 1:
                theta += pi / balls_in_slice
            x, y = slice_radius * cos(theta), slice_radius * sin(theta)
            coords.append((x, y, z))
    return coords

def generate_lissajous(n_balls, radius, t, intro_len):
    coords = []
    phase = (t / 40) * pi / 60 # Sync to ball interval
    for i in xrange(n_balls):
        phase -= pi / 60
        coords.append((radius * sin(7 * phase),
                       radius * cos(4.5 * phase),
                       radius * sin(3 * phase + pi / 4)))
        if phase < 0:
            break
    return coords

N_BALLS = 73
balls_box = generate_box(N_BALLS, 1.5)
print "len(balls_box)", len(balls_box)    
balls_helix = generate_helix(N_BALLS, 1.5, 2.5)
print "len(balls_helix)", len(balls_helix)    
balls_ball = generate_ball(N_BALLS, int(math.sqrt(N_BALLS)), 1.5)
print "len(balls_ball)", len(balls_ball)

sequence = [balls_helix, balls_helix, balls_ball, balls_helix,
            balls_box, balls_box, balls_ball, balls_ball,
            balls_helix, balls_helix, balls_ball, balls_helix,
            balls_ball, balls_ball, balls_box, balls_box]

mspb = 463 # Milliseconds per beat - 129.5 BPM
intro_duration = mspb * 8
switch_duration = 300

def interpolate_ball(t, ball1, ball2):
    k = float(t) / switch_duration
    return [k * x1 + (1-k) * x2 for x1, x2 in zip(ball1, ball2)]

def interpolate_balls(t, balls1, balls2):
    return [interpolate_ball(t, ball1, ball2)
            for ball1, ball2
            in zip(balls1, balls2)]

def get_balls(t):
    if t < intro_duration:
        intro_left = intro_duration - t
        balls = generate_lissajous(N_BALLS, 2.0, t, intro_duration)
        if intro_left < switch_duration:
            balls = interpolate_balls(intro_left, balls, sequence[0])
        return balls
    else:
        t -= intro_duration
        phase = (t / (mspb))
        time_to_next = mspb - (t % mspb)
        if time_to_next < switch_duration:
            return interpolate_balls(time_to_next,
                                     sequence[phase % len(sequence)],
                                     sequence[(phase + 1) % len(sequence)])
        else:
            return sequence[phase % len(sequence)]

def project_to_screen(obj, phi, theta):
    x, y, z = obj
    shift_x, shift_y, shift_z = 0, 0, 0
    scale = 30.
    x, y, z = (x-shift_x)*scale, (y-shift_y)*scale, (z-shift_z)*scale
    
    x, z = (cos(theta)*x - sin(theta)*z, 
            sin(theta)*x + cos(theta)*z)

    x, y = (cos(phi)*x - sin(phi)*y, 
            sin(phi)*x + cos(phi)*y)

    return (x+WIDTH/2, y+HEIGHT/2, z)

# Kludge in correct phase - change this to fit the demo scheduling
demo_location = 151350

def render(time_in_millis):
    #print "In render function, time: " + str(time_in_millis)
    time_in_millis %= demo_location
    
    buf = '\0' * (WIDTH*HEIGHT*4)
    img = pygame.image.frombuffer(buf, (WIDTH, HEIGHT), 'RGBA')
    img.fill((0,0,0,255))
    t = time_in_millis / 1000.
    phi = t * 1.2
    theta = t * 2.528
    
    projected_balls = [project_to_screen(obj, phi, theta)
                       for obj
                       in get_balls(time_in_millis)]

    projected_balls.sort(cmp=lambda a, b: cmp(a[2],b[2]))

    for obj in projected_balls:
        img.blit(ball, (obj[0],obj[1]))
    return buf
