import cv2
import numpy as np
import math
import itertools as it

DENSITY = 5
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)


def analyze_block(frame, sp, ep):
    """ Analyze a block """
    block = np.float32(frame[sp[1]:ep[1], sp[0]:ep[0], 0:3])
    block = block.reshape((-1, 3))
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, .2)
    compactness, labels, palette = cv2.kmeans(block, N_COLORS, None, criteria,
                                              2, cv2.KMEANS_RANDOM_CENTERS)
    unique, counts = np.unique(labels.flatten(), return_counts=True)
    for index, count in zip(unique, counts):
        print(index, count, palette[index])
    return frame


def analyze(frame):
    """ Analyze the full frame """
    height, width, d = frame.shape
    n = DENSITY
    m = math.ceil(n * (height / width))
    dx = width / n
    dy = height / m
    for x, y in 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, sp, ep)
    return frame


if __name__ == '__main__':
    camera = cv2.VideoCapture(0)

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

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

    camera.release()
    cv2.destroyAllWindows()