sk6812-client/snowflakes.py

#!/usr/bin/env python

import sk6812_multistrip as sk6812
import time
import math
import sys
import random

s = sk6812.SK6812(sys.argv[1], 2703)

phase = 0
nled = 16
nstrip = 8

interval = 1.0/60

#scale = 0.15
scale = 0.5

strip = 0

tiltphase = 0

#phase_speeds = [13, 23, 42, 5]
#phase_speeds = [37, 41, 43, 47]
phase_speeds = [701, 709, 719, 727]

class Framebuffer:
	def __init__(self):
		self.clear()

	def clear(self):
		self.r = [0] * (nstrip * nled)
		self.g = [0] * (nstrip * nled)
		self.b = [0] * (nstrip * nled)
		self.w = [0] * (nstrip * nled)

	def setpixel(self, strip, led, r, g, b, w):
		idx = strip * nled + led
		self.r[idx] = r
		self.g[idx] = g
		self.b[idx] = b
		self.w[idx] = w

	def addtopixel(self, strip, led, r, g, b, w):
		idx = strip * nled + led
		self.r[idx] += r
		self.g[idx] += g
		self.b[idx] += b
		self.w[idx] += w

	def render(self):
		for strip in range(nstrip):
			for led in range(nled):
				idx = strip * nled + led
				s.set_color(strip, led,
						self.r[idx],
						self.g[idx],
						self.b[idx],
						self.w[idx])

class Snowflake:
	def __init__(self):
		self.phase = random.random() * 2 * math.pi
		self.frequency = 2*math.pi * 1
		self.speed = random.random() * 0.1 + 0.05

		self.radius = 1.1
		self.radiussq = self.radius**2

		self.pos_strip = self.base_pos_strip = random.randint(0, nstrip-1)
		self.pos_led = nled-1+self.radius

	def move(self):
		self.phase += self.frequency * interval
		self.pos_strip = self.base_pos_strip + 1 * math.sin(self.phase)

		self.pos_led -= 0.1

	def render(self, framebuffer):
		for strip in range(int(math.floor(self.pos_strip-self.radius)), int(math.ceil(self.pos_strip+self.radius))):
			for led in range(int(math.floor(self.pos_led-self.radius)), int(math.ceil(self.pos_led+self.radius))):
				if led < 0 or led >= nled:
					continue

				distsq = (strip - self.pos_strip)**2 + (led - self.pos_led)**2

				brightness = 1 - distsq / self.radiussq
				if brightness < 0:
					continue # brightness 0 -> no effect -> save some calculations

				norm_strip = strip
				while norm_strip < 0:
					norm_strip += nstrip
				while norm_strip >= nstrip:
					norm_strip -= nstrip

				framebuffer.addtopixel(norm_strip, led, 0, 0, int(32*brightness), int(48*brightness))

	def has_fallen(self):
		return self.pos_led < -self.radius

snowflakes = []
framebuffer = Framebuffer()

frame = 0

snowlevel = [0] * nstrip

while True:
	# garbage collection
	if True: #frame % 10 == 0:
		visible_snowflakes = []
		for flake in snowflakes:
			if not flake.has_fallen():
				visible_snowflakes.append(flake)
			else:
				idx = int(flake.base_pos_strip)
				if idx < 0:
					idx += nstrip
				elif idx >= nstrip:
					idx -= nstrip

				snowlevel[idx] += 0.2

		snowflakes = visible_snowflakes

	# spawn new flakes
	if random.random() < 0.1:
		snowflakes.append(Snowflake())
	
	framebuffer.clear()

	# move and render the flakes
	for flake in snowflakes:
		flake.move()
		flake.render(framebuffer)

	# add snow on the ground
	for strip in range(nstrip):
		led = 0
		level = snowlevel[strip]
		while level > 0:
			if level >= 1:
				brightness = 1
			else:
				brightness = level

			framebuffer.addtopixel(strip, led, 0, int(8*brightness), int(32*brightness), int(12*brightness))

			level -= 1
			led += 1

		snowlevel[strip] *= 0.999

	framebuffer.render()

	s.commit()

	time.sleep(interval)
	frame += 1