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

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


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):
    """ Analyze a block """
    block = frame[sp[1]:ep[1], sp[0]:ep[0], 0:3]
    average_color = np.average(block, (0, 1))
    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)
    average_color = [int(x) for x in average_color]
    frame = cv2.rectangle(frame, sp2, ep2, average_color, -1)
    average_color = np.uint8([[average_color]])
    h, s, v = cv2.cvtColor(average_color, cv2.COLOR_BGR2HSV)[0][0]

    # Configure the oscillator
    osc.send_message(
        "/radio",
        [int(index),
         float(h / 255),
         float(s / 255),
         float(v / 255), .5])

    # Draw the image
    for thickness, color in ((3, (0, 0, 0)), (1, (255, 255, 255))):
        frame = cv2.putText(frame,
                            text=f"{h:03d} {s:03d} {v:03d}",
                            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 """
    height, width, d = frame.shape
    n = DENSITY
    m = math.ceil(n * (height / width))
    dx = width / n
    dy = height / m
    for index, (x, y) in enumerate(it.product(range(n), range(m))):
        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)
    return frame


if __name__ == '__main__':
    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()