import gradio as gr
import torch
from PIL import Image
# from transformers import AutoModel, AutoProcessor 

# ==========================================
# 1. Load Model Weights 
# ==========================================
MODEL_ID = "yangxiaoda/SpatialLogic-ckpt-tag"
device = "cuda" if torch.cuda.is_available() else "cpu"

print(f"Loading model to {device}...")
# [REPLACE THIS] with your actual model loading code
# processor = AutoProcessor.from_pretrained(MODEL_ID)
# model = AutoModel.from_pretrained(MODEL_ID).to(device)
print("Model loaded successfully!")


# ==========================================
# 2. Define Core Inference Function
# ==========================================
def compare_images(img1, img2, action_desc):
    """
    img1: PIL.Image
    img2: PIL.Image
    action_desc: string
    """
    # 1. Input Validation
    if img1 is None or img2 is None:
        return "⚠️ Please ensure both images (img1 and img2) are uploaded!"
    if not action_desc.strip():
        return "⚠️ Please enter an action description!"

    try:
        # [REPLACE THIS] with your actual inference code
        # Example pseudo-code:
        # inputs = processor(text=action_desc, images=[img1, img2], return_tensors="pt").to(device)
        # outputs = model(**inputs)
        # score1, score2 = outputs.logits 
        
        # Dummy scores for UI testing (Replace with real logic)
        score_img1 = 0.8  
        score_img2 = 0.4  

        # 2. Logic Judgment
        if score_img1 > score_img2:
            result = "img1"
            explanation = "The model evaluates that **Image 1 (img1)** is closer to the completed state of the described task."
        elif score_img2 > score_img1:
            result = "img2"
            explanation = "The model evaluates that **Image 2 (img2)** is closer to the completed state of the described task."
        else:
            result = "Tie"
            explanation = "The model evaluates both images equally regarding the task completion."

        # 3. Format Output
        return f"### **Output:** {result}\n\n**Explanation:** {explanation}"
    
    except Exception as e:
        return f"❌ An error occurred during inference: {str(e)}"

# ==========================================
# 3. Build Gradio Interface
# ==========================================
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    # Changed Title to the exact text you requested
    gr.Markdown("# 🧠 SpatialLogic Reasoning Demo")
    # Changed sub-title to English
    gr.Markdown("Upload two images and input an action description. The model will predict which image is closer to the completed state of the task.")
    
    with gr.Row():
        with gr.Column():
            # Changed labels to English
            img1_input = gr.Image(type="pil", label="Image 1 (img1)")
        with gr.Column():
            img2_input = gr.Image(type="pil", label="Image 2 (img2)")
            
    # Changed label and placeholder to English
    text_input = gr.Textbox(
        label="Action Description", 
        placeholder="e.g., Put the apple on the table"
    )
    
    # Changed button text to English
    submit_btn = gr.Button("🤖 Evaluate", variant="primary")
    
    # Output area
    output_text = gr.Markdown(label="Result")

    # Bind function to components
    submit_btn.click(
        fn=compare_images, 
        inputs=[img1_input, img2_input, text_input], 
        outputs=output_text
    )

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