import numpy as np
from matplotlib.pyplot import imread

from kmeans_one import kmeans
from kmeans_one import assign_clusters


def rgb_to_h(image):
    h = np.empty(shape=image.shape[0:2])
    maxc = image.max(axis=2)
    minc = image.min(axis=2)

    h[minc == maxc] = 0
    rc = (maxc - image[:, :, 0]) / (maxc - minc)
    gc = (maxc - image[:, :, 1]) / (maxc - minc)
    bc = (maxc - image[:, :, 2]) / (maxc - minc)

    h[:, image[:, :, 0] == maxc] = bc[image[:, :, 0] == maxc] - gc[[image[:, :, 0] == maxc]]
    h[:, image[:, :, 1] == maxc] = 2.0 + rc[image[:, :, 1] == maxc] - bc[image[:, :, 1] == maxc]
    h[:, image[:, :, 2] == maxc] = 4.0 + gc[image[:, :, 2] == maxc] - rc[image[:, :, 2] == maxc]
    h = (h / 6.0) % 1.0
    return h


if __name__ == '__main__':
    in_image_label = "image/in.png"

    clusters_number = 3
    image = imread(in_image_label)
    out_image = rgb_to_h(image)

    clusters = kmeans(out_image[:, :].ravel(), clusters_number)
    out_image[:, :] = assign_clusters(out_image[:, :].ravel(), clusters).reshape(out_image.shape)