import math
import random
import threading
import time
import numpy as np
import picamera
import picamera.array
import neopixel

TAKE_PICTURES = False#True

duration = 3600*168*4

global_fps = 0
global_lasttime = time.time()
global_motionframe = None

global_ledfps = 0

global_loc = 'Z'  # Z, X, Y, or ZX, or ZY, or XY (in betweens)
global_lastloc = 'Z'
global_locchange = time.time()

TAU_IMAGE = 1.     # Filter characteristic time for the image low-pass temporal filter
TAU_LOCCHANGE = 1. # How long to wait before looking whether the location has changed
MOTION_THRES = 1000. # Minimal summed motion to trigger change

# A bunch of constants are hardcoded in the LED routines

# Setting concerning what is being printed on the screen for debugging
MOTION_TO_DECADE_DIV = 3
REPORT_PERIOD = 10

# Size of the motion data array (and part of the array that is actually being considered) XXX depends on the resolution requested from the cam
H = 61
W = 82
startH, endH = 1,None
startW, endW = 1,-1
H = H-1
W = W-2





def asciiart(array):
    '''Print an array heatmap in asciiart.'''
    p = " .:-=+*#%@"
    return '-'*(2*W+2) + '\n' + '\n'.join('|%s|'%''.join(p[c]*2 for c in line) for line in array[1:-1]) + '\n' + '-'*(2*W+2)

def updatescreen(duration):
    '''Print debug data.'''
    start = time.time()
    while time.time()-start<duration:
        if global_motionframe is None:
            continue
        #print(asciiart((global_motionframe/MOTION_TO_DECADE_DIV).astype(int).clip(0,9)), flush=False)
        print('LED FPS %d | Cam FPS %d | Loc %s | Group %s | State %s'%(global_ledfps, global_fps, global_loc, global_subgroup, global_currentstate),
              flush=True)
        time.sleep(REPORT_PERIOD)






#####
#
# Motion Stuff
#
#####




class LowPassDetectMotion(picamera.array.PiMotionAnalysis):
    '''The motion detection routine.''' # XXX it works with global variables
    filtered = None
    def analyze(self, a):
        global global_motionframe, global_fps, global_lasttime, global_loc, global_lastloc, global_locchange
        now = time.time()
        # Image acq and filt
        beta = math.exp(-(now-global_lasttime)/TAU_IMAGE)
        alpha = 1-beta
        a = np.sqrt(a['x'].astype(float)**2+a['y'].astype(float)**2)
        a = a[startH:endH,startW:endW]
        if self.filtered is not None:
            a = alpha*self.filtered + beta*a
        self.filtered = a
        global_motionframe = self.filtered
        global_fps = alpha*global_fps + beta/(now-global_lasttime)
        global_lasttime = now
        # Most active location
        if now-global_locchange>TAU_LOCCHANGE:
            grid = a.reshape(4,H//4,5,W//5).sum(axis=(1,3))
            motion = np.max(grid)
            h,w = [int(_) for _ in  np.unravel_index(np.argmax(grid), grid.shape)]
            if motion>MOTION_THRES and (h,w) not in {(2,0),(3,0),(2,4),(3,4)}:
                if h<2 and w<2:
                    loc = 'X'
                elif h<2 and w>2:
                    loc = 'Y'
                elif h>2 and w>1 and w<4:
                    loc = 'Z'
                else:
                    if w<2:
                        loc = 'ZX'
                    elif w>2:
                        loc = 'ZY'
                    else:
                        loc = 'XY'
                if global_loc != loc:
                    global_locchange = now
                    global_lastloc = global_loc
                    global_loc = loc




#####
#
# LED Stuff
#
#####




LED_PIN        = 18      # GPIO pin connected to the pixels (18 uses PWM!).
LED_FREQ_HZ    = 800000  # LED signal frequency in hertz (usually 800khz)
LED_DMA        = 10      # DMA channel to use for generating signal (try 10)
LED_BRIGHTNESS = 255     # Set to 0 for darkest and 255 for brightest
LED_INVERT     = False   # True to invert the signal (when using NPN transistor level shift)
LED_CHANNEL    = 0       # set to '1' for GPIOs 13, 19, 41, 45 or 53
LED_STRIP = neopixel.ws.WS2811_STRIP_GRB

B = 61
S = 35

LS = [0,B,S,S,S,B,S,S,S,B,S,S,B,B,B,S]

LED_COUNT      = sum(LS) # Number of LED pixels.
strip = neopixel.Adafruit_NeoPixel(LED_COUNT, LED_PIN, LED_FREQ_HZ, LED_DMA, LED_INVERT, LED_BRIGHTNESS, LED_CHANNEL, LED_STRIP)
strip.begin()

cumLS = np.cumsum(LS)

ZZ = [[6,7,8,9]]
XX = [[10,11,12], [15,16]]
YY = [[2,3,4,5]]
XYZ = [[12,13,14,15]]
ZYX = [[1,2],[5,6],[9,10]]

zz = [7]
xy = [8,14]
yx = [6,9]
xx = [11]
zy = [10,16] # and 1 = (16-1)%(len(LS)-1)+1
yz = [12,15]
yy = [3]
zx = [4,13]
xz = [2,5]

I = np.arange(LED_COUNT)

MARGIN = 3

def cycle_to_ledmask(cycle):
    mask = np.zeros_like(I).astype(bool)
    for stretch in cycle:
        for start, end in zip(stretch[:-1],stretch[1:]):
            mask[cumLS[start-1]+MARGIN:cumLS[end-1]-MARGIN] = True
    return mask

SUBGROUPS = {
        'ZZ' : cycle_to_ledmask(ZZ),
        'XX' : cycle_to_ledmask(XX),
        'YY' : cycle_to_ledmask(YY),
        'XYZ': cycle_to_ledmask(XYZ),
        'ZYX': cycle_to_ledmask(ZYX),
        }

def vertice_to_ledmask(vertice):
    mask = np.zeros_like(I).astype(bool)
    for loc in vertice:
        down = cumLS[loc-1]
        up = cumLS[(loc-1)%(len(LS)-1)]
        mask[up:up+MARGIN] = True
        mask[down-MARGIN:down] = True
    return mask

OPERATORS = {
        'zz': vertice_to_ledmask(zz),
        'xy': vertice_to_ledmask(xy),
        'yx': vertice_to_ledmask(yx),
        'xx': vertice_to_ledmask(xx),
        'zy': vertice_to_ledmask(zy),
        'yz': vertice_to_ledmask(yz),
        'yy': vertice_to_ledmask(yy),
        'zx': vertice_to_ledmask(zx),
        'xz': vertice_to_ledmask(xz),
        }

EDGE_MASK = np.logical_or.reduce([v for v in SUBGROUPS.values()])
VERTICE_MASK = np.logical_or.reduce([v for v in OPERATORS.values()])

def show(red,green,blue):
    red   = np.clip(red  ,0,255)
    green = np.clip(green,0,255)
    blue  = np.clip(blue ,0,255)
    colors = red << 16 | green << 8 | blue
    for j,c in enumerate(colors):
        c = int(c)
        strip.setPixelColor(j, c)
    strip.show()

def startupsequence():
    step = 0.2
    brightness = 30
    colors = np.zeros_like(I)
    colors[VERTICE_MASK] = brightness
    show(colors,colors,colors)
    time.sleep(step)
    for name in 'ZZ YY XX XYZ ZYX'.split():
        colors = np.zeros_like(I)
        colors[SUBGROUPS[name]] = brightness
        print(name)
        show(colors,colors,colors)
        time.sleep(step)
    for name in 'zz xy yx xx yz zy yy zx xz'.split():
        colors = np.zeros_like(I)
        colors[OPERATORS[name]] = brightness
        print(name)
        show(colors,colors,colors)
        time.sleep(step)
    for mask in [EDGE_MASK,VERTICE_MASK]*3:
        colors = np.zeros_like(I)
        colors[mask] = brightness//6
        print('-')
        show(colors,colors,colors)
        time.sleep(step)





#####
#
# Putting it together and adding the whole quantum mess
#
#####




LOC_TO_SUBGROUP = {
        'Z':'ZZ',
        'X':'XX',
        'Y':'YY',
        }

TRANSITION_TO_SUBGROUP = {
        ('Z' ,'ZY'):'ZYX',
        ('Z' ,'ZX'):'XYZ',
        ('X' ,'ZY'):'ZYX', # arbitrary
        ('Y' ,'ZX'):'XYZ', # arbitrary
        ('X' ,'XY'):'XYZ',
        ('X' ,'ZX'):'ZYX',
        ('Y' ,'XY'):'ZYX',
        ('Y' ,'ZY'):'XYZ',
        }

SUBGROUP_TO_OPS = {
        'ZZ':  {'zz', 'xy', 'yx'},
        'XX':  {'xx', 'zy', 'yz'},
        'YY':  {'yy', 'xz', 'zx'},
        'ZYX': {'zy', 'yx', 'xz'},
        'XYZ': {'xy', 'yz', 'zx'},
        }

SUBGROUP_TO_PARITY = { # 0 corresponds to 1 and 1 corresponds to -1
        'ZZ': 0,
        'YY': 0,
        'XX': 0,
        'XYZ': 1,
        'ZYX': 1,
        }

def complete_state(state, group, parity):
    if len(state) == 2:
        lastoperator = (group - state.keys()).pop()
        value = (parity+sum(state.values()))%2
        return {lastoperator:value, **state}
    newoperator = (group-state.keys()).pop()
    return complete_state({newoperator:random.randint(0,1), **state}, group, parity)



global_subgroup = LOC_TO_SUBGROUP[global_loc]
global_currentstate = complete_state({}, SUBGROUP_TO_OPS[global_subgroup], SUBGROUP_TO_PARITY[global_subgroup])
def leds(duration):
    ls = [_/2/math.pi for _ in [37, 21]]
    ts = [_/2/math.pi for _ in [7.83,  10.21]]
    p = 50
    brightness = 30
    I = np.arange(strip.numPixels())
    start = time.time()
    c = 0
    while time.time()-start<duration:
        t = time.time()-start
        colors = np.zeros_like(I)

        # Background 
        for L,T in zip(ls,ts):
            colors += (((np.sin(I/L-t/T)+1)/2)**p*brightness).astype(int)
            colors += (((np.sin(I/L+t/T)+1)/2)**p*brightness).astype(int)
        colors[VERTICE_MASK] = 0

        # Additions due to a selected location
        tchange = time.time()-global_locchange
        decay = 1# math.exp(-tchange/20)
        selection_light  = brightness/2 * (1-decay)
        selection_light += brightness*2 * decay * ((-math.cos(tchange * 1/5 * 2*3.14)+1)/2)**0.5
        global global_subgroup, global_currentstate
        old_subgroup = global_subgroup
        if (global_lastloc, global_loc) in TRANSITION_TO_SUBGROUP:
            global_subgroup = TRANSITION_TO_SUBGROUP[(global_lastloc, global_loc)]
        elif global_loc in LOC_TO_SUBGROUP:
            global_subgroup = LOC_TO_SUBGROUP[global_loc]
        decay_out = math.exp(-max(0,tchange-60)/10)
        selection_light *= decay_out
        colors[SUBGROUPS[global_subgroup]] += int(selection_light)
        r,g,b = colors*85//58,colors*72//58,colors*1
        #85  72  58

        # Quantum measurement simulation
        if old_subgroup!=global_subgroup:
            preserved_ops = SUBGROUP_TO_OPS[old_subgroup].intersection(SUBGROUP_TO_OPS[global_subgroup])
            global_currentstate = complete_state({op:global_currentstate[op] for op in preserved_ops},
                                                 SUBGROUP_TO_OPS[global_subgroup],
                                                 SUBGROUP_TO_PARITY[global_subgroup])

        for op, val in global_currentstate.items():
            if val==0:
                b[OPERATORS[op]] += int(brightness*decay_out)
            elif val==1:
                r[OPERATORS[op]] += int(brightness*decay_out)

        show(r,g,b)

        # Performance counters
        global global_ledfps
        c += 1
        global_ledfps = c/t







def capture_from_iter(iterator, duration):
    start = time.time()
    while time.time()-start<duration:
        if TAKE_PICTURES:
            f = next(iterator)
        time.sleep(1)


#startupsequence()


with picamera.PiCamera(sensor_mode=4) as cam:
    with LowPassDetectMotion(cam) as output:
        cam.sensor_mode = 4                    # XXX It is questionable whether this really works properly.
        cam.sensor_mode = 4                    # XXX If the resolution is changed, the W and H variables might need to be modified.
        cam.framerate = 10                     # XXX
        cam.resolution = 1296,972              # XXX
        cam.sensor_mode = 4                    # XXX
        cam.sensor_mode = 4                    # XXX

        cam.start_recording('/dev/null', format='h264', motion_output=output)
        tprint = threading.Thread(target=updatescreen,args=(duration,))
        tleds  = threading.Thread(target=leds,args=(duration,))
        cam_iter = cam.capture_continuous('image{timestamp:%y-%m-%d-%H-%M-%S-%f}.jpg', format='jpeg', use_video_port=False, resize=(1296//2,972//2), quality=30, thumbnail=None, bayer=False)
        tcapture = threading.Thread(target=capture_from_iter,args=(cam_iter,duration))
        tprint.start()
        tleds.start()
        tcapture.start()
        tprint.join()
        tleds.join()
        tcapture.join()
        cam.stop_recording()