import gradio as gr
import torch
from PIL import Image
from transformers import pipeline

# 1. Initialize specialized Hugging Face pipelines
# Using BLIP for structural backdrop captioning (Skyscrapers, Promenades, etc.)
captioner = pipeline("image-text-to-text", model="Salesforce/blip-image-captioning-base")

# Using ViT fine-tuned on iNaturalist for fine-grained vegetation/flora recognition
plant_classifier = pipeline("image-classification", model="microsoft/swin-tiny-patch4-window7-224")

def process_park_landscape(input_image):
    if input_image is None:
        return None, "No image uploaded.", "No image uploaded."
    
    # --- Model 1: Custom YOLOv8 (Placeholder for your amenities/palms weights) ---
    # In practice: 
    # model = ultralytics.YOLO("your_fine_tuned_weights.pt")
    # results = model(input_image)
    # annotated_img = results[0].plot()
    annotated_img = input_image  # Placeholder fallback
    
    # --- Model 2: Structural & Backdrop Captioning ---
    try:
        caption_output = captioner(input_image)
        scene_description = caption_output[0]['generated_text']
    except Exception as e:
        scene_description = f"Error generating scene analysis: {str(e)}"
        
    # --- Model 3: Fine-Grained Species/Flora Identification ---
    try:
        classifications = plant_classifier(input_image)
        # Format top 3 predictions cleanly for gr.Label
        species_predictions = {pred['label']: pred['score'] for pred in classifications[:3]}
    except Exception as e:
        species_predictions = {"Error classifying vegetation": 1.0}

    return annotated_img, species_predictions, scene_description

# --- Gradio UI Custom Styling ---
css = """
.gradio-container { background-color: #fcfbfa; font-family: sans-serif; }
.feedback-header { color: #2e4a36; font-weight: 600; }
button.primary-btn { background-color: #3b5944 !important; color: white !important; }
"""

with gr.Blocks(css=css) as demo:
    gr.Markdown("<h1 style='color: #2e4a36; text-align: center;'>🌴 ADParks12K Ecovision Explorer</h1>")
    gr.Markdown(
        "<p style='text-align: center;'>An interdisciplinary computer vision platform analyzing "
        "the visual ecology, amenities, and urban-nature interfaces of Abu Dhabi's park network.</p>"
    )
    
    gr.HTML("<hr style='border: 0; height: 1px; background: #e0dbd5;'>")
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("<h3 class='feedback-header'>📷 Input Landscape</h3>")
            input_img = gr.Image(type="pil", label="Upload Park Image")
            
            # Setting up examples based on your established park typologies
            gr.Examples(
                examples=[
                    ["Samples/https___lh3.googleusercontent.com_gps-cs-s_AC9h4npGgaSZQuynATiQ0QgXcueek9T8wxKmMbBnvOuc7r72rkiAfxtqoK9tcpfMyl7Oe_v00HUI8ShxjzrXXFZkqFbvyHZzsdk-wcai6EX.jpg"],
                    ["Samples/https___lh3.googleusercontent.com_gps-cs-s_AC9h4nqB53PBGnEtxYZmhBOb3h0w8Ks3_yHjrnriWGBmKwV7k3m3fwIGkymj37uZ6LFYyjJUUNGh0_FNCcOaNIr1XVs_Oqxw5oVV6Gb7-_c.jpg"],
                    ["Samples/https___lh3.googleusercontent.com_gps-cs-s_AC9h4nqBWfCJ47sTdL1mCFClYJQrfk3bWGNCJZfE1uV9ZFLRlTH2wFFD5hHZquU9Bomnfjcdg7bP37SOdDjJcWXurMdl99WgZh2Z04MZuLL.jpg"],
                    ["Samples/https___lh3.googleusercontent.com_gps-cs-s_AC9h4nqDaKD8g5W1Ddsz6Fe4JcONGlHE1JI_mrpk4ftyPnjwvq1JvSzvt5AAy82JnmUSWxS_6f8eCLnd8ZxfkacPrugwxJx3kl85Pbfogf_.jpg"],
                    ["Samples/https___lh3.googleusercontent.com_gps-cs-s_AC9h4nqHLYFCU9Ilg34jGH2cip_fR2_5bRnuunZ9mnQ8_j9F5sgo6Z4Cz7vHWS0X-IY9sjGESSfsO9i0wJFdQSJTOueNE_jUX9Zq2pjfPdw.jpg"],
                    ["Samples/https___lh3.googleusercontent.com_gps-cs-s_AC9h4nqL9m4DvUmXLt--tnlMczmHc9RFi1Xm9sV-RbBBVNFOxo1HCNx9A6uHypRmg7ceVPU6M4bdvOlHiEEhvAbkLxj-R2w9L9n2B_wkI4d.jpg"]
                ],
                inputs=input_img,
                label="Explore Dataset Typology Examples"
            )
            
            submit_btn = gr.Button("Run Ecological Inference", variant="primary", elem_classes=["primary-btn"])
            
        with gr.Column(scale=1):
            gr.Markdown("<h3 class='feedback-header'>🔍 Multi-Model Analytical Breakdown</h3>")
            
            with gr.Tab("Object Detection"):
                out_detection = gr.Image(label="Park Amenities & Infrastructure (YOLO)")
                
            with gr.Tab("Fine-Grained Species"):
                out_species = gr.Label(num_top_classes=3, label="Top Predicted Flora/Fauna Species (iNaturalist ViT)")
                
            with gr.Tab("Urban Context"):
                out_scene = gr.Textbox(label="Backdrop & Setting Description (BLIP)")

    submit_btn.click(
        fn=process_park_landscape,
        inputs=[input_img],
        outputs=[out_detection, out_species, out_scene]
    )

if __name__ == "__main__":
    demo.launch()