import numpy as np
from PIL import Image
from sklearn.cluster import KMeans


def join_images(a: Image.Image, b: Image.Image):
    result = Image.new(a.mode, (a.width + b.width, max(a.height, b.height)))
    result.paste(a)
    result.paste(b, (a.width, 0))

    return result


def generate_palette_image(model: KMeans, size=40):
    image = Image.new("RGB", (0, size))

    for cluster_center in model.cluster_centers_:
        color = tuple(map(int, cluster_center))
        part = Image.new("RGB", (40, 40), color)
        image = join_images(image, part)

    return image


def extract_color_clusters(image_array: np.ndarray):
    w, h, d = image_array.shape
    pixels = image_array.reshape(w * h, d)
    model = KMeans(n_clusters=4).fit(pixels)

    return model