from contextlib import ExitStack

import gradio as gr
import numpy as np

from myapp.colorutils import array_to_hex
from myapp.palette import extract_color_clusters, iter_color_shades

with gr.Blocks() as demo:
    image = gr.Image("static/vulture.webp")
    n_colors = gr.Slider(1, 16, 4, step=1)
    button = gr.Button()

    @gr.render(inputs=[image, n_colors])
    def render_palette(image_array: np.ndarray, n_clusers: int):
        model = extract_color_clusters(image_array, n_clusers)
        cluster_shades = iter_color_shades(model, (0, 0.2, 0.4, 0.6))

        with ExitStack() as stack:
            for i, cluster in enumerate(cluster_shades):
                if i % n_clusers == 0:
                    stack.pop_all().close()
                    stack.enter_context(gr.Group())
                    stack.enter_context(gr.Row(variant="compact"))

                gr.ColorPicker(array_to_hex(cluster), container=False)


if __name__ == "__main__":
    demo.launch()