import math
import time
import itertools as it
import cv2
import colorsys
from pythonosc import udp_client
import numpy as np

DENSITY = 4
RED = [0, 0, 255]
N_COLORS = 3
LAST_MESSAGE_TIME = 0


def draw_rectangle(frame, sp, ep):
    """ Draw a rectangle on the frame """
    return cv2.rectangle(frame, sp, ep, RED, 1)


def analyze_block(frame, osc, index, sp, ep, send=False):
    """ Analyze a block """
    block = frame[sp[1]:ep[1], sp[0]:ep[0], 0:3]
    average_color = np.average(block, (0, 1)) / 255
    h, s, v = colorsys.rgb_to_hsv(*average_color)

    # Configure the oscillator
    if send:
        print("Sending message", h, s, v)
        osc.send_message(
            "/radio",
            [int(index), float(h),
             float(s), float(v), .5])

    # Draw a blob of the average color
    sp2 = tuple(int(x) for x in (2 * np.array(sp) + np.array(ep)) / 3)
    ep2 = tuple(int(x) for x in (np.array(sp) + 2 * np.array(ep)) / 3)
    frame = cv2.rectangle(frame, sp2, ep2,
                          [int(x * 255) for x in average_color], -1)

    # Draw the image
    for thickness, color in ((3, (0, 0, 0)), (1, (255, 255, 255))):
        frame = cv2.putText(frame,
                            text=f"{h:.2f} {s:.2f} {v:.2f}",
                            org=(sp[0] + 10, sp[1] + 20),
                            fontScale=.5,
                            color=color,
                            fontFace=2,
                            thickness=thickness)
    return frame


def analyze(frame, osc):
    """ Analyze the full frame """
    global LAST_MESSAGE_TIME
    height, width, d = frame.shape
    n = DENSITY
    m = math.ceil(n * (height / width))
    dx = width / n
    dy = height / m
    send = False
    if time.time() - LAST_MESSAGE_TIME > 0.1:
        LAST_MESSAGE_TIME = time.time()
        send = True
    for index, (y, x) in enumerate(it.product(range(m), range(n))):
        sp = (int(x * dx), int(y * dy))
        ep = (int(x * dx + dx), int(y * dy + dy))
        frame = draw_rectangle(frame, sp, ep)
        frame = analyze_block(frame, osc, index, sp, ep, send)
    return frame


if __name__ == '__main__':
    camera = cv2.VideoCapture("/dev/video2")
    # camera = cv2.VideoCapture(0)
    osc = udp_client.SimpleUDPClient("0.0.0.0", 5005)

    while True:
        ret, frame = camera.read()
        frame = analyze(frame, osc)
        cv2.imshow('Input', frame)

        c = cv2.waitKey(1)
        if c == 27:
            break

    camera.release()
    cv2.destroyAllWindows()