from tkinter import TOP
import pyautogui
import scipy.cluster as cluster
import scipy
import sys
import numpy as np
from time import sleep, time
import numpy
import signal
import atexit
from time import perf_counter as timestamp
from telemetrix_rpi_pico import telemetrix_rpi_pico
from multiprocessing import Pool
import PIL.ImageGrab

class size:
    width = None
    height = None

def sigint_handler(signal=None, frame=None):
    print ('KeyboardInterrupt is caught')
    board.neopixel_clear()
    sleep(0.75)
    board.shutdown()
    sys.exit(0)

FAST_MODE = False
def find_dorminant_color(im):
    if FAST_MODE:
        color = numpy.reshape(numpy.asarray(im),(-1,3))
        color = numpy.median(color,axis=0)
        
        return color
    else:
        NUM_CLUSTERS = 3

        im = im.resize((50, 50))
        ar = np.asarray(im)
        shape = ar.shape
        ar = ar.reshape(np.product(shape[:2]), shape[2]).astype(float)
        codes, dist = cluster.vq.kmeans(ar, NUM_CLUSTERS)

        vecs, dist = cluster.vq.vq(ar, codes) 
        counts, bins = np.histogram(vecs, len(codes))

        index_max = np.argmax(counts)
        peak = codes[index_max]

        return peak
if __name__ == '__main__':
    signal.signal(signal.SIGINT, sigint_handler)
    atexit.register(sigint_handler)
    global NUM_LEDS,BOARDER_SIZE,TOP_LEDS,RIGHT_LEDS,BUTTOM_LEDS,LEFT_LEDS,screenshot,board

    board = telemetrix_rpi_pico.TelemetrixRpiPico()

    

    BOARDER_SIZE = 150
    TOP_LEDS = 19
    RIGHT_LEDS = 11
    BUTTOM_LEDS = 17
    LEFT_LEDS = 10
    REDUCTION = 4
    NUM_LEDS = TOP_LEDS+RIGHT_LEDS+BUTTOM_LEDS+LEFT_LEDS

    board.set_pin_mode_neopixel(pin_number=2,num_pixels=NUM_LEDS)
    board.neopixel_clear(auto_show=True)

    size = size()

    size.width=pyautogui.size().width/REDUCTION
    size.height=pyautogui.size().height/REDUCTION

    board.neopixel_show()
    screenshot = pyautogui.screenshot()
    pool = Pool(16)
    
    while True:
        try:
            screenshot = PIL.ImageGrab.grab()
            lastTime = timestamp()
            chunk = []
            top = screenshot.crop(box=[0,0,size.width,BOARDER_SIZE])
            left = screenshot.crop(box=[0,0,BOARDER_SIZE,size.height])
            buttom = screenshot.crop(box=[0,size.height-BOARDER_SIZE,size.width,size.height])
            right = screenshot.crop(box=[size.width-BOARDER_SIZE,0,size.width,size.height])
            for i in range(0,TOP_LEDS):
                segment = top.crop(box=[i*size.width/TOP_LEDS,0,(i+1)*size.width/TOP_LEDS,BOARDER_SIZE])
                chunk.insert(i,segment)
            
            for i in range(0,LEFT_LEDS):
                segment = left.crop(box=[0,i*size.height/LEFT_LEDS,BOARDER_SIZE,(i+1)*size.height/LEFT_LEDS])
                chunk.insert(TOP_LEDS+RIGHT_LEDS+BUTTOM_LEDS+LEFT_LEDS-i,segment)
            
            for i in range(0,BUTTOM_LEDS):
                segment = buttom.crop(box=[i*size.width/BUTTOM_LEDS,0,(i+1)*size.width/BUTTOM_LEDS,BOARDER_SIZE])
                chunk.insert(TOP_LEDS+RIGHT_LEDS+BUTTOM_LEDS-i,segment)
            board.neopixel_show()
            
            for i in range(0,RIGHT_LEDS):
                segment = right.crop(box=[0,i*size.height/RIGHT_LEDS,BOARDER_SIZE,(i+1)*size.height/RIGHT_LEDS])
                chunk.insert(i+TOP_LEDS,segment)
            colors = pool.map(find_dorminant_color,chunk)
            for i in range(0,len(colors)):
                board.neo_pixel_set_value(i,r=int(colors[i][0]),g=int(colors[i][1]),b=int(colors[i][2]))
            board.neopixel_show()
        except Exception as e:
            print(e)
            print("Retrying")
        #print("loop time : "+str(timestamp()-lastTime))