from gradio_client import Client
import gradio as gr

# Define paths to reference images for shapes
shape_images = {
    "Circle": "circle.png",
    "Square": "square.png"
    # "Triangle": "triangle.png"
}

# Define color palette as hex codes or color names
color_palette = {
    "Red": "#FF0000",
    "Blue": "#0000FF",
    "Green": "#008000",
    "Yellow": "#FFFF00",
    "Purple": "#800080",
    "Orange": "#FFA500"
}
default_color = "Red"
default_color_hex = color_palette[default_color]

def generate_icon(shape, color, description):
    client = Client("multimodalart/stable-cascade")
    color_hex = color_palette.get(color, color_palette[default_color])
    print(color_hex)
    print(shape)
    print(description)
    # Construct prompt based on user description and selected characteristics
    prompt = (f"Create a {description} icon in the shape of a {shape} with {color} color..")
    
    # Call the API with the constructed prompt
    result = client.predict(
        prompt=prompt,
        negative_prompt="",
        seed=0,
        width=1024,
        height=1024,
        prior_num_inference_steps=20,
        prior_guidance_scale=4,
        decoder_num_inference_steps=10,
        decoder_guidance_scale=0,
        num_images_per_prompt=1,
        api_name="/run"
    )
    
    # Return the result image
    return result

def gradio_interface():
    with gr.Blocks() as demo:
        gr.Markdown("### AI Icon Generator")
        custom_css = """
        <style>
        .custom-button {
            width: 150px; /* Set the width of the button */
            height: 40px; /* Set the height of the button */
            font-size: 16px; /* Set the font size */
            text-align: center; /* Center text */
        }
        .output-image {
            width: 200px; /* Set the width of the output image */
            height: 200px; /* Set the height of the output image */
            object-fit: contain; /* Ensure the image scales proportionally */
        }
        </style>
        """
        gr.HTML(custom_css)
        # Display reference images for shapes as selectable options
        shape_options = gr.Radio(
            choices=list(shape_images.keys()), 
            label="Select Shape"
        )

        # Display shape images with fixed sizes using a layout
        with gr.Row():
            for shape, path in shape_images.items():
                with gr.Column():
                    gr.Image(
                        value=path, 
                        label=shape,
                        type="pil",
                        elem_classes="shape-image",
                        height=100,
                        width=100
                    )  # Remove this line

        # Dropdown for color palette
        color_options = gr.Dropdown(
            choices=list(color_palette.keys()), 
            label="Select Color",
            value=default_color
        )
        # color_hex = gr.Textbox(
        #     label="Color Hex Code", 
        #     value=color_palette[default_color],
        #     visible=False
        # )

        # # Update color hex code when a color is selected
        # def update_color_hex(color):
        #     return color_palette[color]

        # color_options.change(fn=update_color_hex, inputs=[color_options], outputs=[color_hex])

        # Textbox for additional description
        description_text = gr.Textbox(
            label="Describe the Icon", 
            placeholder="e.g., textile, steel, modern, futuristic"
        )

        with gr.Row():
            submit_btn = gr.Button("Generate Icon", elem_classes="custom-button")
            output_image = gr.Image(type="pil", elem_classes="output-image")

        # Connect the function to the button
        submit_btn.click(
            generate_icon,
            inputs=[shape_options, color_options, description_text],
            outputs=[output_image]
        )

    demo.launch(share=True)

# Run the Gradio interface
if __name__ == "__main__":
    gradio_interface()