#!/usr/bin/env python

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

import numpy as np

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

phase = 0
nled = 16
nstrip = 8

interval = 1.0/30

scale = 0.15

strip = 0

tiltphase = 0

ENERGY_MULT = 92
ENERGY_DIV = 100

MAXENERGY_ADD = 160
ENERGY_SUB = 7
ENERGY_PULL_MAX_PCT = 100

HIGHLIGHT_OFFSET = -60

energy = np.zeros( (nled+1, nstrip), dtype=int )
energy_smooth = np.zeros( energy.shape, dtype=int )

COLORMAP_X = [ 0, 70, 120, 200, 225, 255]
COLORMAP_Y = [
        (0, 255,   0),
        (0, 255, 128),
        (0, 255, 192),
        (0, 192, 192),
        (0, 128, 255),
        (0, 128, 300)]


def colormap(x, h):
    i = 0
    while i < len(COLORMAP_X)-1 and COLORMAP_X[i+1] < x:
        i += 1

    if i >= len(COLORMAP_X)-1:
        return hsv2rgbw(*COLORMAP_Y[len(COLORMAP_Y)-1])

    start_x = COLORMAP_X[i]
    end_x = COLORMAP_X[i+1]

    start_col = np.array(COLORMAP_Y[i], dtype=int)
    end_col = np.array(COLORMAP_Y[i+1], dtype=int)

    col = (x - start_x) * 100 // (end_x - start_x) * (end_col - start_col) // 100 + start_col

    col[0] = (col[0] + h + 360) % 360

    return hsv2rgbw(*col)


loop = 0
intensity = MAXENERGY_ADD

while True:

    if loop % 5 == 0:
        intensity = random.randint(MAXENERGY_ADD*3//4, MAXENERGY_ADD*5//4)

    # inject random energy in bottom row
    for i in range(nstrip):
        energy[0, i] += random.randint(0, intensity)

    # pull energy from the cell below
    for led in range(nled, 0, -1):
        for strip in range(nstrip):
            wanted_energy = random.randint(0, ENERGY_PULL_MAX_PCT) * energy[led-1, strip] // 100
            pulled_energy = min(wanted_energy, energy[led-1, strip])
            energy[led, strip] += pulled_energy
            energy[led-1, strip] -= pulled_energy

    # decrease global amount of energy
    for led in range(nled+1):
        for strip in range(nstrip):
            #energy[led, strip] *= ENERGY_MULT
            #energy[led, strip] //= ENERGY_DIV
            energy[led, strip] -= min(ENERGY_SUB, energy[led, strip])

    # smooth the energy distribution
    for led in range(nled):
        for strip in range(nstrip):
            strip_left = (strip + nstrip - 1) % nstrip
            strip_right = (strip + nstrip + 1) % nstrip
            led_above = led + 1
            led_below = led - 1

            energy_smooth[led, strip] = 100 * energy[led, strip]
            energy_smooth[led, strip] += 30 * energy[led, strip_left]
            energy_smooth[led, strip] += 30 * energy[led, strip_right]

            gain = 160

            if led_above < nled:
                energy_smooth[led, strip] += 10 * energy[led_above, strip]
                gain += 10
            if led_below >= 0:
                energy_smooth[led, strip] += 10 * energy[led_below, strip]
                gain += 10

            energy_smooth[led, strip] //= gain

    # update colors
    for led in range(nled):
        for strip in range(nstrip):
            color = colormap(energy_smooth[led, strip], (loop // 20) % 360)**2//255
            s.fade_color(strip, led, color[0], color[1], color[2], color[3])


    s.commit()

    loop += 1

    time.sleep(interval)