import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoImageProcessor, AutoModelForImageClassification
from torch.nn.functional import sigmoid
import torch
from PIL import Image

# Load text emotion model
model_name = "SamLowe/roberta-base-go_emotions"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Load image emotion model
image_model_name = "Celal11/resnet-50-finetuned-FER2013-0.001"
image_processor = AutoImageProcessor.from_pretrained(image_model_name)
image_model = AutoModelForImageClassification.from_pretrained(image_model_name)

# Islamic advice mapping based on emotion
islamic_advice = {
    "sadness": "\"Do not grieve; indeed Allah is with us.\"\n\n**Arabic:** لَا تَحْزَنْ إِنَّ اللَّهَ مَعَنَا\n(Qur'an, At-Tawbah 9:40)",
    "joy": "\"If you are grateful, I will surely increase your favor.\"\n\n**Arabic:** لَئِن شَكَرْتُمْ لَأَزِيدَنَّكُمْ\n(Qur'an, Ibrahim 14:7)",
    "fear": "\"Sufficient for us is Allah, and [He is] the best Disposer of affairs.\"\n\n**Arabic:** حَسْبُنَا اللَّهُ وَنِعْمَ الْوَكِيلُ\n(Qur'an, Al Imran 3:173)",
    "anger": "\"The strong is not the one who overcomes people by his strength, but the strong is the one who controls himself while in anger.\"\n\n(Hadith, Sahih al-Bukhari 6114)",
    "confusion": "\"Seek help through patience and prayer.\"\n\n**Arabic:** وَاسْتَعِينُوا بِالصَّبْرِ وَالصَّلَاةِ\n(Qur'an, Al-Baqarah 2:45)",
    "love": "\"Indeed, those who have believed and done righteous deeds - the Most Merciful will appoint for them affection.\"\n\n**Arabic:** سَيَجْعَلُ لَهُمُ الرَّحْمَٰنُ وُدًّا\n(Qur'an, Maryam 19:96)",
    "neutral": "May Allah always guide your heart in every situation."
}

# Analyze image emotion using processor and model
def analyze_image_emotion(image):
    if image is None:
        return "No image provided.", ""
    inputs = image_processor(images=image, return_tensors="pt")
    with torch.no_grad():
        logits = image_model(**inputs).logits
    probs = torch.nn.functional.softmax(logits, dim=1)[0]
    pred_idx = torch.argmax(probs).item()
    label = image_model.config.id2label[pred_idx].lower()
    score = probs[pred_idx].item()
    islamic = islamic_advice.get(label, "May Allah always guide your heart.")
    return f"{label.capitalize()} ({score:.2f})", islamic

# Emotion label to icon mapping (subset)
emotion_icons = {
    "admiration": "😍",
    "amusement": "😅",
    "anger": "😡",
    "annoyance": "😑",
    "approval": "👍",
    "caring": "💗",
    "confusion": "🤔",
    "curiosity": "😮",
    "desire": "🤤",
    "disappointment": "😞",
    "disapproval": "👎",
    "disgust": "🤮",
    "embarrassment": "😳",
    "excitement": "🎉",
    "fear": "😱",
    "gratitude": "🙏",
    "grief": "😭",
    "joy": "😃",
    "love": "❤️",
    "nervousness": "🤧",
    "optimism": "😊",
    "pride": "😎",
    "realization": "🤯",
    "relief": "😌",
    "remorse": "😔",
    "sadness": "😢",
    "surprise": "😲",
    "neutral": "😐"
}

# Analyze text emotion
def get_emotions(text, threshold):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        logits = model(**inputs).logits
    probs = sigmoid(logits)[0]

    top_idx = torch.argmax(probs).item()
    top_label = model.config.id2label[top_idx].lower()
    top_score = probs[top_idx].item()
    islamic = islamic_advice.get(top_label, "May Allah always guide your heart.")

    labels = [model.config.id2label[i] for i, p in enumerate(probs) if p > threshold]
    icons = [emotion_icons.get(label, '') + ' ' + label.capitalize() + f" ({probs[i]:.2f})" for i, label in enumerate(labels)]

    result = ", ".join(icons) if icons else "No strong emotion detected."
    return result, islamic

# Combined analysis
def analyze_combined(text, threshold, image):
    text_result, text_tip = get_emotions(text, threshold)
    image_result, image_tip = analyze_image_emotion(image)
    return text_result, text_tip, image_result, image_tip

# Gradio UI
custom_css = """
body {
  background: linear-gradient(to right, #f9f9f9, #d4ecff);
  font-family: 'Segoe UI', sans-serif;
}
.gr-button {
  background-color: #007BFF !important;
  color: white !important;
  border-radius: 8px !important;
  font-weight: bold;
}
.gr-button:hover {
  background-color: #0056b3 !important;
}
.gr-textbox {
  border-radius: 8px !important;
  border: 1px solid #ccc !important;
  padding: 10px !important;
}
.output-textbox {
  font-size: 1.5rem;
  font-weight: bold;
  color: #333;
  background-color: #f1f9ff;
  border-radius: 8px;
  padding: 10px;
  border: 1px solid #007BFF;
}
"""

demo = gr.Interface(
    fn=analyze_combined,
    inputs=[
        gr.Textbox(lines=5, placeholder="Write a sentence or a full paragraph...", label="Your Text"),
        gr.Slider(minimum=0.1, maximum=0.9, value=0.3, step=0.05, label="Threshold"),
        gr.Image(type="pil", label="Upload Face Photo")
    ],
    outputs=[
        gr.Textbox(label="Detected Text Emotions", elem_classes=["output-textbox"]),
        gr.Textbox(label="Quranic/Hadith Advice (Text)", elem_classes=["output-textbox"]),
        gr.Textbox(label="Detected Photo Emotion", elem_classes=["output-textbox"]),
        gr.Textbox(label="Quranic/Hadith Advice (Image)", elem_classes=["output-textbox"])
    ],
    title="🥰 Multi-Modal Emotion Detector with Islamic Insight",
    description="Analyze emotion from both text and a facial photo. Then receive inspirational Islamic advice based on your mood.",
    theme="default",
    css=custom_css
)

demo.launch()