blup/generators/unrath.py

#!/usr/bin/env python3

import sys
import random
import argparse
import blup.frame
import blup.output
import blup.animation
import blup.writebml
import colorsys


WIDTH = 22
HEIGHT = 16
DEPTH = 256


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


def get_random_color():
    if random.random() < 0.5:
        s = 1
        v = random.random()
    else:
        s = random.random()
        v = 1
    return colorsys.hsv_to_rgb(random.random(), s, v)


class Unrath:
    def __init__(self, world):
        self.world = world
        self.shape = set([(0, 0)])
        shapesize = random.randint(2, 5)
        while len(self.shape) < shapesize:
            p = random.sample(self.shape, 1)[0]
            self.shape.add((p[0]+random.randint(-1, 1),
                            p[1]+random.randint(-1, 1)))
        minx = min(map(lambda p: p[0], self.shape))
        miny = min(map(lambda p: p[1], self.shape))
        self.shape = {(p[0]+abs(minx), p[1]+abs(miny)) for p in self.shape}
        self.shapewidth = max(map(lambda p: p[0], self.shape))
        shapeheight = max(map(lambda p: p[1], self.shape))

        y = random.randint(0, HEIGHT-shapeheight)
        if random.random() > 0.5:
            d = 1
            x = -self.shapewidth
        else:
            d = -1
            x = WIDTH + self.shapewidth
        self.pos = (x, y)
        self.finished = False
        self.speed = d * random.randint(2, WIDTH/2)
        self.color = get_random_color()

    def update(self, dt, nocheck=False):
        if (self.pos[0] + self.shapewidth < 0 and self.speed < 0) or \
           (self.pos[0] > WIDTH and self.speed > 0):
           self.finished = True
           return
        newpos = (self.pos[0] + dt*self.speed, self.pos[1])
        newpoints = self.get_points(newpos)
        collision = False
        for unrath in self.world.unraethe - {self}:
            if len(unrath.get_points().intersection(newpoints)) > 0:
                collision = True
                if abs(unrath.speed) > abs(self.speed):
                    fast = unrath
                    slow = self
                else:
                    slow = unrath
                    fast = self
                slow.speed = fast.speed * (1.1 + random.random()/10)
                slow.update(dt, True)
        if not collision:
            self.pos = newpos

    def get_points(self, pos=None):
        if pos is None:
            pos = self.pos
        return {(int(pos[0] + s[0]), int(pos[1] + s[1])) for s in self.shape}

    def draw(self, frame):
        for p in self.get_points():
            try:
                frame.setPixel(p[0], p[1], convert_color(self.color))
            except ValueError:
                pass


class World:
    def __init__(self, num_unraethe):
        self.num_unraethe = num_unraethe
        self.unraethe = {}

    def update(self, dt):
        self.unraethe = set(filter(lambda u: not u.finished, self.unraethe))
        if len(self.unraethe) < self.num_unraethe:
            self.unraethe.add(Unrath(self))

        for u in self.unraethe:
            u.update(dt)

    def draw(self, frame):
        for u in self.unraethe:
            u.draw(frame)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Generate Unrath animations')
    parser.add_argument('-d', '--delay', type=int, default=50)
    parser.add_argument('-t', '--time', type=int, default=15)
    parser.add_argument('-n', '--num-unraethe', type=int, default=5)
    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)

    t = 0
    w = World(args.num_unraethe)
    while t < args.time:
        t += args.delay / 1000
        frame = blup.animation.AnimationFrame(dim, args.delay)
        w.update(args.delay / 1000)
        w.draw(frame)
        anim.addFrame(frame)

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