import streamlit as st
from transformers import BlipProcessor, BlipForQuestionAnswering
from PIL import Image
import torch

# Define model names
original_model_name = "Salesforce/blip-vqa-base"
finetuned_model_name = "manik063/blip-mmsd-sarcasm-no-reasoning"

@st.cache_resource
def load_models():
    # Load both processors and models
    processor_orig = BlipProcessor.from_pretrained(original_model_name)
    model_orig = BlipForQuestionAnswering.from_pretrained(original_model_name)
    model_orig = model_orig.to('cuda' if torch.cuda.is_available() else 'cpu')
    model_orig.eval()

    processor_finetuned = BlipProcessor.from_pretrained(finetuned_model_name)
    model_finetuned = BlipForQuestionAnswering.from_pretrained(finetuned_model_name)
    model_finetuned = model_finetuned.to('cuda' if torch.cuda.is_available() else 'cpu')
    model_finetuned.eval()

    return (processor_orig, model_orig), (processor_finetuned, model_finetuned)

# Load models
(orig_processor, orig_model), (finetuned_processor, finetuned_model) = load_models()

st.title("🖼️ Sarcasm Detection: Compare Original vs Fine-tuned BLIP Together")
st.write("Upload an image and caption, and we will check sarcasm with BOTH models.")

uploaded_file = st.file_uploader("Upload Image", type=["png", "jpg", "jpeg"])
caption = st.text_input("Enter Caption:")

if uploaded_file and caption:
    if st.button("Detect Sarcasm 🚀"):
        image = Image.open(uploaded_file).convert("RGB")

        results = {}

        for model_name, processor, model in [
            ("Original BLIP (Salesforce)", orig_processor, orig_model),
            ("Fine-tuned BLIP (Manik063)", finetuned_processor, finetuned_model)
        ]:
            prompt = f"Is this caption sarcastic? Answer strictly 'yes' or 'no': {caption}"
            inputs = processor(images=image, text=prompt, return_tensors="pt").to(model.device)

            with torch.no_grad():
                outputs = model.generate(**inputs)

            generated_text = processor.batch_decode(outputs, skip_special_tokens=True)[0].lower()
            results[model_name] = generated_text

        st.image(image, caption="Uploaded Image", use_column_width=True)
        st.write(f"**Caption:** {caption}")

        # Display individual model outputs
        # for model_name, output in results.items():
        #     st.write(f"**{model_name} Output:** {output}")

        for model_name, output in results.items():
            if isinstance(output, str):
                output = output.lower()
                print(output)
                print(type(output))
                if "yes" in output or "1" in output:
                    st.write(f"**{model_name} Output:** 🔴 Sarcasm Detected!")
                elif "no" in output or "0" in output:
                    st.write(f"**{model_name} Output:** 🟢 Not Sarcastic.")
                else:
                    st.write(f"**{model_name} Output:** ⚠️ Could not determine.")
            else:
                st.write(f"**{model_name} Output:** ⚠️ Invalid output format.")

        # Final decision logic
        # all_outputs = " ".join(results.values())
        # if "yes" in all_outputs:
        #     st.success("🔴 Sarcasm Detected by at least one model!")
        # elif "no" in all_outputs:
        #     st.success("🟢 Not Sarcastic according to both models!")
        # else:
        #     st.warning("⚠️ Models could not decide.")