from collections import namedtuple
from time import sleep
import badge
import ugfx

class Attractor(object):
    
    def __init__(self,dt,scales,variables=['x','y','z'],initial=0.1):
        self.dt = dt
        self.var_count = range(len(variables))
        self.variables = variables
        self.point = namedtuple('point',variables)
        self.scale = namedtuple('scale',['mi','mx'])
        self.scales = [self.scale(*i) for i in scales]
        self.ranges = [1.0 / (s.mx - s.mi) for s in self.scales ]
        self.initial = initial
 
    def delta(self,i,p,dp):
        l = list(p)
        l[i] += dp
        return self.point(*l)
        
    def rk4(self,p):
        k1 = [self.dt*f(p) for f in self.partials]
        k2 = [self.dt*self.partials[i](self.delta(i,p,0.5 * k1[i])) for i in self.var_count]
        k3 = [self.dt*self.partials[i](self.delta(i,p,0.5 * k2[i])) for i in self.var_count]
        k4 = [self.dt*self.partials[i](self.delta(i,p,k3[i])) for i in self.var_count]
        return self.point(*[p[i] + (0.5 * k1[i] + k2[i] + k3[i] + 0.5 * k4[i]) / 3.0 for i in self.var_count])
    
    def euler(self,p):
        return self.point(*[r + self.dt * f(p) for r,f in zip(p,self.partials)])
    
    def trace(self,use_rk4=True):
        p = self.point(*[self.initial for i in self.var_count])
        if use_rk4:
            solver = self.rk4
        else:
            solver = self.euler
        while True:
            new_p = solver(p)
            trigger = new_p.x * p.x < 0.1
            yield tuple([self.ranges[i] * (p[i] - self.scales[i].mi) for i in self.var_count] + [trigger] )
            p = new_p
            
class Lorenz(Attractor):
    
    def __init__(self,sigma=10.0,rho=28.0,beta=8.0/3.0,dt=0.01):
        self.sigma = sigma
        self.rho = rho
        self.beta = beta
        
        self.partials = (lambda p: self.sigma * (p.y - p.x),
            lambda p: p.x * (self.rho - p.z) - p.y, lambda p: p.x * p.y - self.beta * p.z)
                
        super(Lorenz,self).__init__(dt, [(-20,25), (-40,40), (0,60)])

class Chua(Attractor):
    
    def __init__(self,alpha=15.6,beta=25,dt=0.015):
        self.alpha = alpha
        self.beta = beta
        self.partials = (lambda p: self.alpha * (p.y - p.x - self.f(p.x)),
            lambda p: p.x - p.y + p.z, lambda p: -self.beta * p.y)
        
        super(Chua,self).__init__(dt, [(-3,3), (-0.5,0.5), (-5,5)])
        
    def f(self,h):
        return -0.7142857142857143 * h + (-0.21428571428571425)*(abs(h+1)-abs(h-1))

class Rossler(Attractor):
    
    def __init__(self,a=0.1,b=0.1,c=14.0,dt=0.05):
        self.a = a
        self.b = b
        self.c = c
        self.partials = (lambda p: -p.y - p.z, lambda p: p.x + self.a * p.y,
            lambda p: self.b + p.z * (p.x - self.c))
                    
        super(Rossler,self).__init__(dt, [(-20,25), (-25,20), (-5,110)],initial=1.0)
        
class App(object):
    
    def __init__(self):
        self.WIDTH = 296
        self.HEIGHT  = 128
        self.scale = 0
        self.animal = 0
        self.zoo = [Lorenz,Chua,Rossler]
        self.reset = False
        self.toggle = False
        self.blinken = False
        self.scalers = (lambda x,y,z: [int(x*self.WIDTH),self.HEIGHT-int(z*self.HEIGHT)],
            lambda x,y,z: [int(x*self.WIDTH),self.HEIGHT-int(y*self.HEIGHT)],
            lambda x,y,z: [int(y*self.WIDTH),self.HEIGHT-int(z*self.HEIGHT)])
        badge.init()
        ugfx.init()
        ugfx.input_init()
        self.reset_trace()
        self.running = True
        ugfx.input_attach(ugfx.BTN_B, self.quit)
        ugfx.input_attach(ugfx.BTN_A, self.blinken_lichten)
        ugfx.input_attach(ugfx.JOY_LEFT, self.prev_scale)
        ugfx.input_attach(ugfx.JOY_RIGHT, self.next_scale)
        ugfx.input_attach(ugfx.JOY_DOWN, self.prev_animal)
        ugfx.input_attach(ugfx.JOY_UP, self.next_animal)
        
    def stream(self):
        powers = [2**i for i in range(6)]
        while True:
            for i in range(64):
                green = [127 if p & i else 0 for p in powers]
                lights = []
                # Wot no itertools.chain?
                for g in green:
                    lights.append(g)
                    lights.append(0)
                    lights.append(0)
                    lights.append(0)
                print(lights)
                yield bytes(lights)
        
    def reset_trace(self):
        ugfx.clear(ugfx.WHITE)
        self.attract = self.zoo[self.animal]()
        self.scaled_trace = (self.scalers[self.scale](x,y,z) for x,y,z,tr in self.attract.trace())
        self.points = [self.scaled_trace.__next__() for i in range(128)]
        for p1, p2 in zip(self.points,self.points[1:]):
            ugfx.line(* p1 + p2 + [ugfx.BLACK])
    
    
    def blinken_lichten(self,junk):
        if not self.toggle:
            self.toggle = True
            self.blinken = not self.blinken
    
    def quit(self,junk):
        self.running = False
    
    def prev_scale(self,junk):
        if not self.reset:
            self.scale = (self.scale - 1 ) % 3
            self.reset = True
            ugfx.flush()
        
    def next_scale(self,junk):
        if not self.reset:
            self.scale = (self.scale + 1 ) % 3
            self.reset = True
            ugfx.flush()

    def prev_animal(self,junk):
        if not self.reset:
            self.animal = (self.animal - 1 ) % 3
            self.reset = True
            ugfx.flush()
        
    def next_animal(self,junk):
        if not self.reset:
            self.animal = (self.animal + 1 ) % 3
            self.reset = True    
            ugfx.flush()
        
    def run(self):
        lights = self.stream()
        while self.running:
            if self.toggle:
                if self.blinken:
                    badge.leds_init()
                else:
                    badge.leds_disable()
                self.toggle = False
            if self.reset:
                ugfx.flush()
                self.reset_trace()
                self.reset = False
            if self.blinken:
                badge.leds_send_data(lights.__next__(), 24)
            ugfx.line(* self.points[0] + self.points[1] + [ugfx.WHITE])
            del self.points[0]
            self.points.append(self.scaled_trace.__next__())
            ugfx.line(* self.points[-2] + self.points[-1] + [ugfx.BLACK])
            ugfx.flush()
        
app = App()
app.run()