import leds
import display
import utime
from math import sqrt, atan, pi

num_leds = 15
rocket_colours = [[0, 0, 255],
                    [255, 128, 0],
                    [0, 255, 0]]

def led_pos(i):
    if i < 11:
        return [int(round(126.5-7.75*i)),0]
    elif i < 15:
        return [[79,0],[79,159],[159,0],[0,0]][i-11]

def avg(l, x=0):
    def _avg(l):
        if isinstance(l,(int, float)):
            return l
        else:
            return sum([_avg(i) for i in l])/len(l)
    if x:
        return round(_avg(l), x)
    else:
        return int(round(_avg(l), 0))


def get_col(x,y, phase=0, hsv=False):
    def hsv_to_rgb(h, s, v):
            if s == 0.0: return (v, v, v)
            i = int(h*6.) # XXX assume int() truncates!
            f = (h*6.)-i
            p,q,t = v*(1.-s), v*(1.-s*f), v*(1.-s*(1.-f))
            i%=6
            if i == 0: return (v, t, p)
            if i == 1: return (q, v, p)
            if i == 2: return (p, v, t)
            if i == 3: return (p, q, v)
            if i == 4: return (t, p, v)
            if i == 5: return (v, p, q)

    def transform(x,y):
        return [x-79.5, y-39.5]

    def arg(x,y):
        if x > 0:
            return atan(y/x)/(2*pi) + int(y < 0)
        elif x < 0:
            return atan(y/x)/(2*pi) + 1/2
        else:
            return (1/(2*pi) + int(y < 0))/2

    if hsv:
        return [360*((arg(*transform(x,y)) + phase) % 1), 1, 0.5]
    else:
        return [int(round(255 * i)) for i in hsv_to_rgb((arg(*transform(x,y)) + phase) % 1, 1, 0.5)]

def set_display(disp, phase):
    def mute(x):
        pass

    for y in range(80):
        [None for x in range(160) if mute(disp.pixel(x,y, col=get_col(x,y, phase=phase)))]
        disp.update()



def set_colours(phase, colours):
    leds.set_all_hsv(colours[:15])

    with display.open() as disp:
        set_display(disp, phase)
        disp.backlight(31)
        disp.update()

def get_rb(n):
    i = 0
    while True:
        phase = i/num_leds
        yield [phase, [get_col(*led_pos(k), phase=phase, hsv=True) for k in range(n)]]
        i = (i+1)%n


with display.open() as disp:
    disp.clear()

while True:
    [set_colours(c, colours) for c, colours in get_rb(num_leds)]