from colors import *
import math
import numpy as np
import random
import screen
from screen import Screen

all_anims = {}

def anim(name):
    def _register(c):
        all_anims[name] = c
        return c
    return _register

def get(current, aid):
    bits = aid.split('/')
    if len(bits) != 3:
        return None
    typ = all_anims.get(bits[1], None)
    bright = int(bits[2], 16)
    if type(current) == typ:
        current.brightness = bright
        return current
    return typ(brightness=bright)

def encode(anim_name, brightness):
    return "anim/{}/{:02x}".format(anim_name, brightness)

class Anim:
    desired_frame_time = 0.02

    def __init__(self, brightness=255):
        self.brightness = brightness

    def get_screen(self):
        self.step()
        return Screen(
            colordata=[[self.get_color(x, y) for x in range(Screen.W)]
                       for y in range(Screen.H)],
            brightness=self.brightness)

    def step(self):
        pass

    def get_color(self, x, y):
        pass

@anim("AJB")
class AJB(Anim):
    desired_frame_time = 0.03
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.h = 0

    def step(self):
        self.h += self.desired_frame_time / 2

    def get_color(self, x, y):
        s = 0.3
        v = 0.6
        if random.random() < 0.05:
            s = 0.1
            v = 1.0
        return filter_hsv((self.h + 0.05 * x, s, v))

@anim("JDB")
class JDB(Anim):
    desired_frame_time = 0.03
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def get_color(self, x, y):
        h = 0.6
        s = 0.25
        v = 0.5
        if random.random() < 0.05:
            s = 0.1
            v = 1.0
        return filter_hsv((h, s, v))

@anim("KCS")
class KCS(Anim):
    desired_frame_time = 0.03
    hstep = 0.03
    ha = 0.17
    hb = 0.45
    s = 0.3

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.t = 0
        self.inc = True

    def step(self):
        if self.inc:
            self.t += self.hstep
            if self.t > 1:
                self.t = 1
                self.inc = False
        else:
            self.t -= self.hstep
            if self.t < 0:
                self.t = 0
                self.inc = True
        h = self.ha + self.t * (self.hb - self.ha)
        self.c1 = filter_hsv((h, self.s, 1.0))
        h = self.ha + (1 - self.t) * (self.hb - self.ha)
        self.c2 = filter_hsv((h, self.s, 1.0))

    def get_color(self, x, y):
        if (x % 2 == 0) ^ (y % 2 == 0):
            #return (255, 0, 0)
            return self.c1
        return self.c2

@anim("XMS")
class XMS(Anim):
    desired_frame_time = 0.01
    spot_period = 8.0
    r = (0.0, 1.0, 1.0)
    gh = 0.30

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.y = 0
        self.spot = None
        self.spot_age = 0
        self.spot_outer = self.r
        self.spot_inner = self.r

    def step(self):
        self.spot_age += self.desired_frame_time
        if self.spot_age > self.spot_period or self.spot is None:
            self.spot = (
                random.choice(list(range(Screen.W))),
                random.choice(list(range(Screen.H))))
            self.spot_age = 0
        if self.spot_age < 0.2 * self.spot_period:
            t = self.spot_age / (0.2 * self.spot_period)
            self.spot_inner = (0.0, 1.0 - t, 1.0)
            self.spot_outer = self.spot_inner
        elif self.spot_age < 0.4 * self.spot_period:
            t = self.spot_age - 0.2 * self.spot_period
            t /= 0.2 * self.spot_period
            self.spot_inner = (0.0, 0.0, 1.0)
            self.spot_outer = (self.gh, t, 1.0)
        elif self.spot_age < 0.6 * self.spot_period:
            self.spot_inner = (0.0, 0.0, 1.0)
            self.spot_outer = (self.gh, 1.0, 1.0)
        elif self.spot_age < 0.8 * self.spot_period:
            t = self.spot_age - 0.6 * self.spot_period
            t /= 0.2 * self.spot_period
            t = 1. - t
            self.spot_inner = (0.0, 0.0, 1.0)
            self.spot_outer = (self.gh, t, 1.0)
        else:
            t = self.spot_age - 0.8 * self.spot_period
            t /= (0.2 * self.spot_period)
            t = 1. - t
            self.spot_inner = (0.0, 1.0 - t, 1.0)
            self.spot_outer = self.spot_inner

    def get_color(self, x, y):
        c = self.r
        spot_dist = abs(x - self.spot[0]) + abs(y - self.spot[1])
        if spot_dist <= 1:
            c = self.spot_inner
        elif spot_dist <= 3:
            c = self.spot_outer
        return self.drop_sat(c)

    def drop_sat(self, hsv):
        h, s, v = hsv
        s = pow(s, 4)
        return filter_hsv((h, s, v))

@anim("NYE")
class NYE(Anim):
    desired_frame_time = 0.01
    center = np.array(((Screen.W - 1) / 2., (Screen.H - 1) / 2.), dtype=np.float64)
    width = 3
    speed = 4

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.roff = 0

    def step(self):
        self.roff += self.speed * self.desired_frame_time

    def get_color(self, x, y):
        v = np.array((x, y), dtype=np.float64)
        d = (np.linalg.norm(self.center - v) - self.roff) % self.width / self.width
        if d < 0.25:
            s = 0.0
        elif d < 0.5:
            s = pow((d - 0.25) / 0.25, 4)
        elif d < 0.75:
            s = 1.0
        else:
            s = pow(1. - (d - 0.75) / 0.25, 4)
        return filter_hsv((0.7, s, 1.0))

@anim("IDL")
class IDL(Anim):
    desired_frame_time = 0.01
    vec = np.array((2, 5), dtype=np.float64)
    vec /= np.linalg.norm(vec)
    w = 45
    speed = 2.8
    base_h = 0.85
    h_wid = 0.25

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.off = 0

    def step(self):
        self.off += self.speed * self.desired_frame_time

    def get_color(self, x, y):
        v = np.array((x, y))
        d = (v.dot(self.vec) + self.off) % self.w / self.w
        if d < 0.5:
            t = d * 2
        else:
            t = 1. - (d - 0.5) * 2.
        return filter_hsv((self.base_h + self.h_wid * t, 0.7, 1.0))