import gradio as gr
import cv2
import numpy as np

# from panorama import stitch_images


def process_images(image_files, crop, y_start, y_end, x_start, x_end):
    if not image_files:
        return "No images uploaded.", None

    # Extract the first element of each tuple if it's a tuple
    extracted_images = []
    for item in image_files:
        if isinstance(item, tuple):
            # If item is (img, None), extract 'img' only
            img = item[0]
        else:
            # If for some reason it's not a tuple, assume it's already the image
            img = item

        # Check if img is not None and is a numpy array
        if img is not None and isinstance(img, np.ndarray):
            extracted_images.append(img)

    # If after extraction no valid images found
    if not extracted_images:
        return "No valid images found for stitching.", None

    try:
        # Define crop coordinates
        if crop:
            crop_coords = (int(y_start), int(y_end), int(x_start), int(x_end))
        else:
            crop_coords = None

        # Initialize OpenCV's Stitcher
        stitcher = cv2.Stitcher_create()

        # Perform stitching
        status, panorama = stitcher.stitch(extracted_images)

        if status != cv2.Stitcher_OK:
            raise RuntimeError(f"Stitching failed with status code {status}.")

        if crop:
            y_start, y_end, x_start, x_end = crop_coords
            cropped_panorama = panorama[y_start:y_end, x_start:x_end]
            return panorama, cropped_panorama

        return panorama, None

    except Exception as e:
        # Raise a gradio Error to avoid returning image outputs as text
        raise gr.Error(f"Error during stitching: {str(e)}")


# Define Gradio interface using Blocks for better layout control
with gr.Blocks(title="Creating Panorama using OpenCV") as demo:
    gr.Markdown("# Creating Panorama using OpenCV")
    gr.Markdown("Upload the series of images sequentially and get the perfect Panorama!")

    with gr.Row(equal_height=True):
        
        examples = [
            "./scene/scene1.jpg",
            "./scene/scene2.jpg",
            "./scene/scene3.jpg",
            "./scene/scene4.jpg",
            "./scene/scene5.jpg",
            "./scene/scene6.jpg",
            "./scene/scene7.jpg",
            "./scene/scene8.jpg",
        ]
        with gr.Column():
            image_upload = gr.Gallery(
                value=examples,
                columns=4,
                label="Upload Images",
                file_types=["image"],
                format="jpeg",
                type="numpy",
            )
            crop_checkbox = gr.Checkbox(label="Apply Cropping to Panorama", value=False)
            with gr.Row():
                y_start = gr.Number(label="Crop Y Start", value=90, interactive=True)
                y_end = gr.Number(label="Crop Y End", value=867, interactive=True)
            with gr.Row():
                x_start = gr.Number(label="Crop X Start", value=1, interactive=True)
                x_end = gr.Number(label="Crop X End", value=2000, interactive=True)
            stitch_button = gr.Button("Stitch Images")
    with gr.Row():
        with gr.Column():
            stitched_output = gr.Image(label="Stitched Panorama")
            cropped_output = gr.Image(label="Cropped Panorama")


    # Disable cropping input fields if checkbox is not selected
    def toggle_crop_inputs(crop):
        return gr.update(visible=crop), gr.update(visible=crop), gr.update(visible=crop), gr.update(visible=crop)

    crop_checkbox.change(fn=toggle_crop_inputs, inputs=[crop_checkbox], outputs=[y_start, y_end, x_start, x_end])

    # Define the button click event
    stitch_button.click(
        fn=process_images,
        inputs=[image_upload, crop_checkbox, y_start, y_end, x_start, x_end],
        outputs=[stitched_output, cropped_output],
    )

    gr.Markdown(
        """
        **Note**:
        - Ensure that the uploaded images have overlapping regions for successful stitching.
        - Cropping coordinates should be within the dimensions of the stitched panorama.
        """
    )

# Launch the Gradio interface
demo.launch(show_error=True)