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

from src.pandora_removal import PandoraRemoval, PandoraConfig


# Initialize the model
@torch.no_grad()
def initialize_model():
    """Initialize PANDORA model with default config."""
    config = PandoraConfig(
        device="cuda"  # Force GPU usage, no CPU fallback
    )
    model = PandoraRemoval(config=config)
    model.load_model()
    return model


# Global model instance
print("Loading PANDORA model...")
model = initialize_model()
print("✓ Model loaded successfully!")


def process_image(
    image_dict: dict,
    border_size: int,
    guidance_scale: float,
    num_steps: int,
    seed: int
) -> list[Image.Image]:
    """Process image with mask using PANDORA.

    Args:
        image_dict: Dictionary from ImageEditor with 'background' and 'layers'
        border_size: Border dilation size
        guidance_scale: Guidance scale for generation
        num_steps: Number of diffusion steps
        seed: Random seed for reproducibility

    Returns:
        List containing output image with object removed
    """
    if image_dict is None:
        raise ValueError("Please upload an image")

    # Extract source image
    source_image = image_dict.get('background')
    if source_image is None:
        raise ValueError("No image found")

    # Handle both PIL Image and numpy array
    if isinstance(source_image, np.ndarray):
        source_image = Image.fromarray(source_image).convert("RGB")
    elif isinstance(source_image, Image.Image):
        source_image = source_image.convert("RGB")
    else:
        raise ValueError(f"Unexpected image type: {type(source_image)}")

    # Extract mask from layers
    layers = image_dict.get('layers', [])
    if not layers or len(layers) == 0:
        raise ValueError("Please draw a mask on the object you want to remove")

    # Get the first layer (mask)
    mask_array = layers[0]
    if isinstance(mask_array, np.ndarray):
        mask_image = Image.fromarray(mask_array).convert('L')
    elif isinstance(mask_array, Image.Image):
        mask_image = mask_array.convert('L')
    else:
        raise ValueError(f"Unexpected mask type: {type(mask_array)}")

    # Convert mask to tensor
    from torchvision import transforms
    mask_tensor = transforms.ToTensor()(mask_image).to(model.device)

    # Set seed if specified
    if seed != -1:
        torch.manual_seed(seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed(seed)

    # Process with PANDORA
    result = model.remove_object(
        image=source_image,
        mask=mask_tensor,
        border_size=border_size,
        guidance_scale=guidance_scale,
        num_steps=num_steps
    )

    return [result]


# Create Gradio interface with MimicBrush-style layout
with gr.Blocks(title="PANDORA Object Removal") as demo:
    with gr.Row():
        with gr.Column():
            gr.Markdown("# PANDORA: Zero-shot Object Removal via Diffusion Models")

            with gr.Row():
                output_gallery = gr.Gallery(
                    label='Output',
                    show_label=True,
                    elem_id="gallery",
                    columns=1,
                    height=768
                )

            with gr.Accordion("Advanced Options", open=True):
                border_size = gr.Slider(
                    label="Border Size",
                    minimum=0,
                    maximum=50,
                    value=17,
                    step=1,
                    info="Size of border around mask"
                )
                num_steps = gr.Slider(
                    label="Steps",
                    minimum=20,
                    maximum=100,
                    value=50,
                    step=1
                )
                guidance_scale = gr.Slider(
                    label="Guidance Scale (LADG)",
                    minimum=0.0,
                    maximum=15.0,
                    value=1.3,  # Optimal value from PANDORACode config
                    step=0.1,
                    info="Lower values (1.0-2.0) work best for object removal"
                )
                seed = gr.Slider(
                    label="Seed",
                    minimum=-1,
                    maximum=999999999,
                    step=1,
                    value=-1,
                    info="Set to -1 for random"
                )

            gr.Markdown("### Tutorial")
            gr.Markdown("1. Upload an image")
            gr.Markdown("2. Use the draw button to mask the object you want to remove")
            gr.Markdown("3. Adjust parameters if needed")
            gr.Markdown("4. Click 'Remove Object' to generate")

            gr.Markdown("### Tips")
            gr.Markdown("- Draw the mask slightly larger than the object")
            gr.Markdown("- Border size: 2-5 for small objects, 15-20 for medium, 20-30 for large objects")
            gr.Markdown("- Guidance Scale (LADG): Keep around 1.0-2.0 for best results (default 1.3)")
            gr.Markdown("- Images are automatically resized to 768x768 for processing")

        with gr.Column():
            gr.Markdown("# Upload your image")
            gr.Markdown("### Tips: You can adjust the brush size")

            input_image = gr.ImageEditor(
                label="Source Image",
                type="pil",
                brush=gr.Brush(colors=["#FFFFFF"], default_size=30, color_mode="fixed"),
                layers=False,
                interactive=True,
                height=512
            )

            run_button = gr.Button(value="Remove Object", variant="primary", size="lg")

    # Examples section (add example images to demo_examples/ directory)
    with gr.Row():
        gr.Examples(
            examples=[
                ["./demo_examples/img_1_original.png", 17, 7.5, 50, -1],
                ["./demo_examples/img_2_original.png", 17, 7.5, 50, -1],
                ["./demo_examples/img_3_original.png", 17, 7.5, 50, -1],
                ["./demo_examples/img_4_original.png", 17, 7.5, 50, -1],
            ],
            inputs=[
                input_image,
                border_size,
                guidance_scale,
                num_steps,
                seed
            ],
            cache_examples=False,
            examples_per_page=10
        )

    # Footer
    gr.Markdown(
        """
        ---
        ### About PANDORA

        PANDORA is a state-of-the-art object removal method using diffusion models with attention control.
        It intelligently removes objects from images while preserving the background and generating plausible content.

        **Key Features:**
        - Zero-shot object removal (no training required)
        - Attention-controlled diffusion
        - High-quality background reconstruction
        - Flexible control parameters

        **Citation:** If you use this work, please cite our paper.
        """
    )

    # Event handler
    run_button.click(
        fn=process_image,
        inputs=[
            input_image,
            border_size,
            guidance_scale,
            num_steps,
            seed
        ],
        outputs=output_gallery
    )


if __name__ == "__main__":
    demo.launch(
        share=True,
        server_name="0.0.0.0",
        server_port=7860
    )