blup/generators/orbs.py

#!/usr/bin/env python3

import sys
import random
import numpy as np
import argparse
import blup.frame
import blup.output
import blup.animation
import blup.writebml
import colorsys


WIDTH = 22
HEIGHT = 16
DEPTH = 256


def fill_trace(points):
    for i in range(len(points)):
        yield points[i]
        if i == len(points)-1:
            break
        v = points[i+1] - points[i]
        l = int(np.linalg.norm(v))
        for j in range(l):
            yield points[i] + v * (j+1)/(l+1)


def convert_color(c):
    return list(map(lambda x: int(round(x*(DEPTH-1))), c))


def get_random_color():
    return colorsys.hsv_to_rgb(random.random(), 1, 1)


class Orb:
    def __init__(self, pos, speed, tracelen, color, tracecolor, boxsize):
        self.pos = pos
        self.speed = speed
        self.tracelen = tracelen
        self.color = color
        self.tracecolor = tracecolor
        self.boxsize = boxsize
        self.trace = []

    def update(self, dt):
        newpos = self.pos + self.speed * dt
        #newpos = self.pos
        for dim in [0, 1]:
            if newpos[dim] >= self.boxsize[dim] or newpos[dim] <= 0:
                newpos[dim] = 0 if newpos[dim]<0 else self.boxsize[dim]-1
                self.speed[dim] *= -1
        self.speed += np.array([random.randint(-5, 6), random.randint(-5, 6)])
        self.trace.insert(0, self.pos)
        self.trace = self.trace[:self.tracelen]
        self.pos = newpos

    def draw(self, frame):
        for i, p in enumerate(fill_trace([self.pos] + self.trace)):
            f = (self.tracelen - i) / self.tracelen
            try:
                c = colorsys.rgb_to_hsv(*self.tracecolor)
                s = 0.1 if random.random()>0.9 else c[1]
                c = colorsys.hsv_to_rgb(c[0], s, c[2]*f)
                frame.setPixel(int(p[0]), int(p[1]), convert_color(c))
            except ValueError:
                pass

        offs = [ np.array([x,y]) for x in [-1,0,1] for y in [-1,0,1] ]
        points = []
        for o in offs:
            p = self.pos.astype('int') + o
            if random.random() > 0.2 and not (p == np.array([0,0])).all():
                points.append(p)
        for p in points:
            try:
                c = colorsys.rgb_to_hsv(*self.color)
                c = colorsys.hsv_to_rgb(c[0], c[1], c[2]*0.2)
                frame.setPixel(p[0], p[1], convert_color(c))
            except ValueError:
                pass

        frame.setPixel(int(self.pos[0]), int(self.pos[1]), convert_color(self.color))


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Generate plasma animations')
    parser.add_argument('-d', '--delay', type=int, default=50)
    parser.add_argument('-t', '--time', type=int, default=15)
    parser.add_argument('-n', '--num-orbs', type=int, default=2)
    parser.add_argument('output_file')
    args = parser.parse_args()

    dim = blup.frame.FrameDimension(WIDTH, HEIGHT, DEPTH, 3)
    anim = blup.animation.Animation(dim)
    anim.tags['description'] = ' '.join(sys.argv)

    orbs = []
    for i in range(args.num_orbs):
        orbs.append(Orb(
            np.array([random.randint(0, WIDTH-1),random.randint(0, HEIGHT-1)]),
            np.array([8,10]),
            20,
            get_random_color(),
            get_random_color(),
            np.array([WIDTH, HEIGHT])
        ))

    t = 0
    while t < args.time:
        t += args.delay / 1000
        frame = blup.animation.AnimationFrame(dim, args.delay)

        for orb in orbs:
            orb.update(args.delay / 1000)
            orb.draw(frame)

        anim.addFrame(frame)

    blup.writebml.writeBml(anim, args.output_file)