import os
from typing import Tuple

import gradio as gr
import numpy as np
import pandas as pd
from PIL import Image
from sklearn.cluster import KMeans


def _image_resize(image: Image.Image, pixels: int = 90000, **kwargs):
    rt = (image.size[0] * image.size[1] / pixels) ** 0.5
    if rt > 1.0:
        small_image = image.resize((int(image.size[0] / rt), int(image.size[1] / rt)), **kwargs)
    else:
        small_image = image.copy()
    return small_image


def get_main_colors(image: Image.Image, n: int = 28, pixels: int = 90000) \
        -> Tuple[Image.Image, np.ndarray, np.ndarray, np.ndarray]:
    image = image.copy()
    if image.mode != 'RGB':
        image = image.convert('RGB')
    small_image = _image_resize(image, pixels)

    few_raw = np.asarray(small_image).reshape(-1, 3)
    kmeans = KMeans(n_clusters=n)
    kmeans.fit(few_raw)

    width, height = image.size
    raw = np.asarray(image).reshape(-1, 3)
    colors = kmeans.cluster_centers_.round().astype(np.uint8)
    prediction = kmeans.predict(raw)
    new_data = colors[prediction].reshape((height, width, 3))
    new_image = Image.fromarray(new_data, mode='RGB')

    cids, counts = np.unique(prediction, return_counts=True)
    counts = np.asarray(list(map(lambda x: x[1], sorted(zip(cids, counts)))))

    return new_image, colors, counts, prediction.reshape((height, width))


def main_func(image: Image.Image, n: int, pixels: int, fixed_width: bool, width: int):
    if fixed_width:
        _width, _height = image.size
        r = width / _width
        new_width, new_height = int(round(_width * r)), int(round(_height * r))
        image = image.resize((new_width, new_height))

    new_image, colors, counts, predictions = get_main_colors(image, n, pixels)

    table = pd.DataFrame({
        'r': colors[:, 0],
        'g': colors[:, 1],
        'b': colors[:, 2],
        'count': counts,
    })
    table['ratio'] = table['count'] / table['count'].sum()
    hexes = []
    for r, g, b in zip(table['r'], table['g'], table['b']):
        hexes.append(f'#{r:02x}{g:02x}{b:02x}')
    table['hex'] = hexes

    new_table = pd.DataFrame({
        'Hex': table['hex'],
        'Pixels': table['count'],
        'Ratio': table['ratio'],
        'Red': table['r'],
        'Green': table['g'],
        'Blue': table['b'],
    }).sort_values('Pixels', ascending=False)

    return new_image, new_table


if __name__ == '__main__':
    pd.set_option("display.precision", 3)
    with gr.Blocks() as demo:
        with gr.Row():
            with gr.Column():
                ch_image = gr.Image(type='pil', label='Original Image')
                with gr.Row():
                    ch_clusters = gr.Slider(value=8, minimum=2, maximum=256, step=2, label='Clusters')
                    ch_pixels = gr.Slider(value=100000, minimum=10000, maximum=1000000, step=10000,
                                          label='Pixels for Clustering')
                    ch_fixed_width = gr.Checkbox(value=True, label='Width Fixed')
                    ch_width = gr.Slider(value=200, minimum=12, maximum=2048, label='Width')

                ch_submit = gr.Button(value='Submit', variant='primary')

            with gr.Column():
                with gr.Tabs():
                    with gr.Tab('Output Image'):
                        ch_output = gr.Image(type='pil', label='Output Image')
                    with gr.Tab('Color Map'):
                        ch_color_map = gr.Dataframe(
                            headers=['Hex', 'Pixels', 'Ratio', 'Red', 'Green', 'Blue'],
                            label='Color Map'
                        )

        ch_submit.click(
            main_func,
            inputs=[ch_image, ch_clusters, ch_pixels, ch_fixed_width, ch_width],
            outputs=[ch_output, ch_color_map],
        )

    demo.queue(os.cpu_count()).launch()