import gradio as gr
from PIL import Image
import os

from image_segmentation_mask_rcnn import segment_person
from insert_person_into_stereo import insert_person_from_combined_stereo
from create_anaglyph import create_anaglyph

# Predefined sample files (make sure these exist in your project directory)
DEFAULT_BACKGROUNDS = {
    "Backyard": "Side_By_Side_3D_Images/sbs_backyard.png",
    "Campus": "Side_By_Side_3D_Images/sbs_campus.png",
    "Downtown": "Side_By_Side_3D_Images/sbs_downtown.png",
    "NEU": "Side_By_Side_3D_Images/sbs_neu.png",
    "STEAM_CLOCK": "Side_By_Side_3D_Images/sbs_steam_clock.png",
    "Trail": "Side_By_Side_3D_Images/sbs_trail.png"
}

DEFAULT_PEOPLE = {
    "PERSON1": "person/person1.jpg",
    "PERSON2": "person/person2.png",
    "PERSON3": "person/person3.png",
    "PERSON4": "person/person4.png",
    "PERSON5": "person/person5.png",
}

def pipeline(person_image, stereo_image, depth, x, y):
    segmented = segment_person(person_image)

    left_image, right_image, _ = insert_person_from_combined_stereo(
        stereo_image=stereo_image,
        segmented_person=segmented,
        depth=depth,
        position=(x, y)
    )

    anaglyph = create_anaglyph(left_image, right_image)
    return anaglyph

def get_image_dimensions(stereo_image, person_image):
    if stereo_image is None or person_image is None:
        return gr.update(), gr.update()

    w_bg, h_bg = stereo_image.size
    w_p, h_p = person_image.size

    max_x = max(10, w_bg // 2 - w_p // 2)  # Ensure > 0
    max_y = max(10, h_bg)  # Ensure > 0

    return gr.update(minimum=0, maximum=max_x, value=max_x // 2), gr.update(minimum=0, maximum=max_y, value=int(h_bg * 0.9))

def main():
    with gr.Blocks() as demo:
        gr.Markdown("# 🧍‍➡️ 3D Anaglyph Composer")

        with gr.Row():
            person_input = gr.Image(type="pil", label="Person Image")
            stereo_input = gr.Image(type="pil", label="Stereo Background")

        # Sample selectors
        with gr.Row():
            with gr.Column():
                gr.Markdown("### Sample People")
                for label, path in DEFAULT_PEOPLE.items():
                    with gr.Row():
                        preview = gr.Image(value=path, label=label, interactive=False, show_label=False, width=128, height=128)
                        use_btn = gr.Button(f"Use {label}")
                        use_btn.click(lambda p=path: Image.open(p), outputs=person_input)

            with gr.Column():
                gr.Markdown("### Sample Backgrounds")
                for label, path in DEFAULT_BACKGROUNDS.items():
                    with gr.Row():
                        preview = gr.Image(value=path, label=label, interactive=False, show_label=False)
                        use_btn = gr.Button(f"Use {label}")
                        use_btn.click(lambda p=path: Image.open(p), outputs=stereo_input)

        depth_input = gr.Dropdown(["close", "medium", "far"], value="medium", label="Depth")
        x_slider = gr.Slider(0, 2000, value=1000, label="Person X Position")
        y_slider = gr.Slider(0, 2000, value=500, label="Person Y Position")

        generate_btn = gr.Button("Generate 3D Anaglyph")
        output_img = gr.Image(type="pil", label="Anaglyph 3D Image")

        # Dynamically update position sliders when images are uploaded
        person_input.change(get_image_dimensions, inputs=[stereo_input, person_input], outputs=[x_slider, y_slider])
        stereo_input.change(get_image_dimensions, inputs=[stereo_input, person_input], outputs=[x_slider, y_slider])

        generate_btn.click(
            fn=pipeline,
            inputs=[person_input, stereo_input, depth_input, x_slider, y_slider],
            outputs=output_img
        )

    demo.launch()

if __name__ == "__main__":
    main()