import os
import gradio as gr
import numpy as np
from PIL import Image

# Placeholder for your actual model
def generate_image(image_idx: int, x: float, y: float) -> Image.Image:
    """
    Replace this with your actual model inference
    """
    # This is just a placeholder - replace with your model
    # Creating a simple gradient image as example output
    width, height = 256, 256
    gradient = np.zeros((height, width, 3), dtype=np.uint8)
    gradient[:, :, 0] = np.linspace(0, 255 * x, width)
    gradient[:, :, 1] = np.linspace(0, 255 * y, height)[:, np.newaxis]
    gradient[:, :, 2] = image_idx * 30  # vary blue channel based on selected image
    return Image.fromarray(gradient)

def process_click(image_idx: int, evt: gr.SelectData) -> Image.Image:
    """
    Process the click event on the coordinate selector
    """
    # Extract coordinates from click event
    x, y = evt.index[0], evt.index[1]
    # Normalize coordinates to [0, 1]
    x, y = x/100, y/100
    # Generate image using the model
    return generate_image(image_idx, x, y)

with gr.Blocks() as demo:
    gr.Markdown("""
    # Interactive Image Generation
    Choose a reference image and click on the coordinate selector to generate a new image.
    """)
    
    with gr.Row():
        # Left column: Reference images and coordinate selector
        with gr.Column(scale=1):
            # Radio buttons for image selection
            image_idx = gr.Radio(
                choices=[i for i in range(4)],  # Replace with your actual number of images
                value=0,
                label="Select Reference Image",
                type="index"
            )
            
            # Display reference images
            gallery = gr.Gallery(
                value=[
                    "image_0.jpg",
                    "image_0.jpg",
                    "image_0.jpg",
                    "image_0.jpg",
                ],
                columns=2,
                rows=2,
                height=300,
                label="Reference Images"
            )
            
            # Coordinate selector (displayed as heatmap for click interaction)
            coord_selector = gr.Plot(
                value=None,
                label="Click to select (x, y) coordinates"
            )
            
            # Initialize the coordinate selector
            def create_selector():
                import plotly.graph_objects as go
                fig = go.Figure()
                
                # Add a square shape
                fig.add_trace(go.Scatter(
                    x=[0, 100, 100, 0, 0],
                    y=[0, 0, 100, 100, 0],
                    mode='lines',
                    line=dict(color='black'),
                    showlegend=False
                ))
                
                # Update layout
                fig.update_layout(
                    width=300,
                    height=300,
                    margin=dict(l=0, r=0, t=0, b=0),
                    xaxis=dict(
                        range=[-5, 105],
                        showgrid=False,
                        zeroline=False,
                        visible=False
                    ),
                    yaxis=dict(
                        range=[-5, 105],
                        showgrid=False,
                        zeroline=False,
                        visible=False,
                        scaleanchor='x'
                    ),
                    plot_bgcolor='white'
                )
                return fig
            
            # Initialize the coordinate selector
            coord_selector.value = create_selector()
        
        # Right column: Generated image
        with gr.Column(scale=1):
            output_image = gr.Image(
                label="Generated Output",
                height=300
            )
    
    # Handle click events
    coord_selector.select(
        process_click,
        inputs=[image_idx],
        outputs=output_image
    )

# Launch the app
demo.launch()