#!/usr/bin/env python3

# inspired by:
# http://lodev.org/cgtutor/plasma.html

import sys
import time
import math
import colorsys
import blup.frame
import blup.output
import blup.animation
import blup.writebml
import argparse

WIDTH = 22
HEIGHT = 16
DEPTH = 256


class Plasma(object):
    def __init__(self, width, height, test_palette=False, zoom=1):
        self.width = width
        self.height = height
        self.test_palette = test_palette
        self.zoom = zoom

        self.pixelvalues = [[0] * height for i in range(width)]
        self.pixelsReady = False
        self.offset = 0
        self.update()

    @property
    def minValue(self):
        return min(list(map(min, self.pixelvalues)))

    @property
    def maxValue(self):
        return max(list(map(max, self.pixelvalues)))

    def applyPalette(self, palette):
        min = self.minValue
        max = self.maxValue
        norm = lambda x: ((x - min) / max)

        pixeldata = [ [None] * self.height for i in range(self.width) ]

        for x in range(self.width):
            rowvalues = self.pixelvalues[x]
            for y in range(self.height):
                if y == 0 and self.test_palette:
                    pixeldata[x][y] = palette.getColorValue(x/(self.width-1))
                else:
                    pixeldata[x][y] = palette.getColorValue(norm(rowvalues[y]))

        return pixeldata

    def update(self):
        #if self.pixelsReady:
        #    return
        self.offset += 1

        for x in range(self.width):
            for y in range(self.height):
                xx = x+self.offset
                yy = y+self.offset
                p = 128 + 128 * math.sin(xx / 9.0 * self.zoom)
                p += 128 + 128 * math.sin(yy / 16.0 * self.zoom)
                p += 128 + 128 * math.sin((yy + xx) / 16.0 * self.zoom)
                p += 128 + 128 * math.sin(math.sqrt(yy*yy + xx*xx * self.zoom) / 16.0)
                p += 128 + 8 * math.sin(math.sqrt(yy* xx * self.zoom) / 16.0)
                self.pixelvalues[x][y] = p
        self.pixelsReady = True

class HSVPalette(object):
    def __init__(self, depth):
        self.offset = 0
        self.cache = {}
        self.depth = depth

    def getColorValue(self, x):
        x = round(x, 3)
        if x in self.cache:
            return self.cache[x]
        c = colorsys.hsv_to_rgb((x + self.offset) % 1, 1, 1)
        c = list(map(lambda x: int(x * (self.depth - 1)), c))
        self.cache[x] = c
        return c

    def update(self):
        self.cache = {}
        self.offset += 0.01
        if self.offset > 1:
            self.offset -= 1

class ExponentialHSVPalette(object):
    def __init__(self, depth):
        self.offset = 0
        self.cache = {}
        self.depth = depth

    def getColorValue(self, x):
        x = round(x, 3)
        if x in self.cache:
            return self.cache[x]
        xx = (x + self.offset) % 1
        a = 4
        b = 3
        c = colorsys.hsv_to_rgb((a**(xx*b)-1)/(a**b), 1, 1)
        c = list(map(lambda x: int(x * (self.depth - 1)), c))
        self.cache[x] = c
        return c

    def update(self):
        self.cache = {}
        self.offset += 0.01
        if self.offset > 1:
            self.offset -= 1

class SingleColorPalette(object):
    def __init__(self, depth):
        self.cache = {}
        self.depth = depth
        self.h = 0

    def getColorValue(self, x):
        x = round(x, 3)
        if x in self.cache:
            return self.cache[x]
        a = 4
        b = 3
        c = colorsys.hsv_to_rgb(self.h, 1, (a**(x*b)-1)/(a**b))
        c = list(map(lambda x: int(x * (self.depth - 1)), c))
        self.cache[x] = c
        return c

    def update(self):
        self.cache = {}
        self.h += 0.005
        if self.h > 1:
            self.h -= 1

##############################################

PALETTES = {
    'hsv': HSVPalette,
    'exponentialhsv': ExponentialHSVPalette,
    'singlecolor': SingleColorPalette,
}

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Generate plasma animations')
    parser.add_argument('-d', '--delay', type=int, default=50)
    parser.add_argument('-n', '--num-frames', type=int, default=200)
    parser.add_argument('-p', '--palette', choices=PALETTES.keys(), default='hsv')
    parser.add_argument('-z', '--zoom', type=float, default=1)
    parser.add_argument('--test-palette', action='store_true')
    parser.add_argument('output_file')
    args = parser.parse_args()

    plasma = Plasma(WIDTH, HEIGHT, args.test_palette, zoom=args.zoom)
    dim = blup.frame.FrameDimension(WIDTH, HEIGHT, DEPTH, 3)
    palette = PALETTES[args.palette](DEPTH)

    anim = blup.animation.Animation(dim)
    anim.tags['description'] = ' '.join(sys.argv)


    for i in range(args.num_frames):
        palette.update()
        plasma.update()
        pixeldata = plasma.applyPalette(palette)
        frame = blup.animation.AnimationFrame(dim, args.delay)

        for x in range(WIDTH):
            for y in range(HEIGHT):
                frame.setPixel(x, y, pixeldata[x][y])
        anim.addFrame(frame)

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