import numpy as np
import spaces
import gradio as gr
from gradio.themes.base import Base
from gradio.themes.utils import colors, fonts, sizes
from typing import Iterable
from PIL import Image
from parallax import generate_parallax_images


class Theme(Base):
    def __init__(
        self,
        *,
        primary_hue: colors.Color | str = colors.neutral,
        secondary_hue: colors.Color | str = colors.neutral,
        neutral_hue: colors.Color | str = colors.neutral,
        spacing_size: sizes.Size | str = sizes.spacing_md,
        radius_size: sizes.Size | str = sizes.radius_md,
        text_size: sizes.Size | str = sizes.text_md,
        font: fonts.Font | str | Iterable[fonts.Font | str] = (fonts.GoogleFont('Barlow'), 'ui-sans-serif', 'sans-serif'),
        font_mono: fonts.Font| str | Iterable[fonts.Font | str] = (fonts.GoogleFont('IBM Plex Mono'), 'ui-monospace', 'monospace',),
    ):
        super().__init__(
            primary_hue=primary_hue,
            secondary_hue=secondary_hue,
            neutral_hue=neutral_hue,
            spacing_size=spacing_size,
            radius_size=radius_size,
            text_size=text_size,
            font=font,
            font_mono=font_mono,
        )
        super().set(
            color_accent='rgb(0 231 255 / 1)',
            slider_color='rgb(0 231 255 / 1)',
            slider_color_dark='rgb(0 231 255 / 1)',
            button_primary_background_fill='rgb(0 231 255 / 1)',
            button_primary_background_fill_hover='rgb(0 231 255 / .75)',
            button_primary_text_color='#ffffff',
            button_primary_background_fill_dark='rgb(0 231 255 / 1)',
            button_primary_background_fill_hover_dark='rgb(0 231 255 / .75)',
            button_primary_text_color_dark='#ffffff',
            loader_color='rgb(255 199 229 / 1)',
            loader_color_dark='rgb(255 199 229 / 1)'
        )


@spaces.GPU(duration=30)
def generate(image: np.ndarray, n: int):
    return generate_parallax_images(Image.fromarray(image, 'RGBA'), n)


with gr.Blocks(theme=Theme()) as demo:
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(image_mode='RGBA', label='Input')
            input_layers = gr.Number(value=5, precision=0, minimum=2, maximum=10, step=1, label='Layers')
            generate_button = gr.Button('Generate')
        with gr.Column():
            output_gallery = gr.Gallery(label='Outputs', columns=5)

    generate_button.click(fn=generate, inputs=[input_image, input_layers], outputs=output_gallery, api_name='generate')


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