blup/generators/fireworks.py

#!/usr/bin/env python3

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

WIDTH = 22
HEIGHT = 16
DEPTH = 256

G = np.array([0, 9.81])

class Fireball:
    def __init__(self, startpos, startspeed, startcolor, lifetime, t0):
        self.startpos, self.startspeed, self.startcolor, self.lifetime = \
            startpos, startspeed, startcolor, lifetime
        self.t0 = t0
        self.pos = startpos
        self.color = startcolor
        self.dead = False

    def update_pos(self, t):
        t = t - self.t0
        if t > self.lifetime:
            self.dead = True
            return
        self.pos = 0.5 * G * t**2 + self.startspeed * t + self.startpos
        self.color = self.startcolor * (1 - t/self.lifetime)
        if random.randint(1, 10) > 8:
            self.color = np.array([0, 0, 0])

    def draw(self, frame):
        if not self.dead:
            frame.setPixel(int(round(self.pos[0])), int(round(self.pos[1])),
                           tuple(map(lambda x: int(round(x)), self.color)))

class Rocket:
    def __init__(self, startpos, startspeed, explode_height, t0, color):
        self.startpos, self.startspeed, self.explode_height, self.t0 = \
            startpos, startspeed, explode_height, t0
        self.color = color
        self.pos = None
        self.exploded = False

    def update(self, t):
        t = t - self.t0
        if t < 0 or self.exploded:
            return
        self.pos = 0.5 * G * t**2 + self.startspeed * t + self.startpos
        if self.pos[1] <= self.explode_height:
            self.exploded = True
            self.color = np.array([0, 0, 0])

    def draw(self, frame):
        if self.pos is None or self.exploded:
            return
        frame.setPixel(int(round(self.pos[0])), int(round(self.pos[1])),
                       tuple(map(lambda x: int(round(x)), self.color)))

        


def getRandomColor(maxval):
    return list(map(lambda x: int(round(x*maxval)), colorsys.hsv_to_rgb(random.random(), 1, 1)))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Generate plasma animations')
    parser.add_argument('-d', '--delay', type=int, default=50)
    parser.add_argument('-n', '--num-rockets', type=int, default=3)
    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)

    rockets = []
    for i in range(args.num_rockets):
        rockets.append(
            Rocket(
                np.array([random.randint(0, dim.width-1), dim.height-1]),
                np.array([random.randint(-8, 8), -20 + random.randint(-2, 2)]),
                4 + random.randint(-2, 2),
                random.randint(0, 40) / 10,
                np.array([dim.depth-1] * 3)
            )
        )

    fbs = []
    t = 0
    while True:
        frame = blup.animation.AnimationFrame(dim, args.delay)

        for fb in fbs:
            fb.update_pos(t)
            try:
                fb.draw(frame)
            except ValueError:
                continue
        fbs = list(filter(lambda f: not f.dead, fbs))

        for r in rockets:
            r.update(t)
            try:
                r.draw(frame)
            except ValueError:
                continue
            if r.exploded:
                color = np.array(getRandomColor(dim.depth-1))
                for i in range(random.randint(15, 25)):
                    v = np.array([
                        random.choice([-1, 1]) * random.randint(0, 25),
                        random.choice([-1, 1]) * random.randint(0, 25),
                    ])
                    fbs.append(Fireball(r.pos, v, color, 1.5 + random.randint(-1, 1), t))
        rockets = list(filter(lambda r: not r.exploded, rockets))

        anim.addFrame(frame)

        if len(rockets) == 0 and len(fbs) == 0:
            break
        t += args.delay / 1000

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