import gradio as gr
import joblib
from huggingface_hub import hf_hub_download
import json

# ================== YOUR MODEL REPO ==================
REPO_ID = "ShanukaB/English_LR_Model_New"

# Download model files from Hugging Face Hub
tfidf_path = hf_hub_download(repo_id=REPO_ID, filename="tfidf_vectorizer.joblib")
model_path = hf_hub_download(repo_id=REPO_ID, filename="logreg_model.joblib")
map_path   = hf_hub_download(repo_id=REPO_ID, filename="emotion_map.json")

# Load the files
tfidf = joblib.load(tfidf_path)
model = joblib.load(model_path)

with open(map_path, "r") as f:
    emotion_map = json.load(f)

# Emotion names for nice display
emotion_display = {
    0: "sadness",
    1: "joy",
    2: "love",
    3: "anger",
    4: "fear",
    5: "surprise"
}

def predict_emotion(text):
    if not text or not text.strip():
        return "Please write something 😶"
    
    # Transform and predict
    text_tfidf = tfidf.transform([text])
    pred_idx = model.predict(text_tfidf)[0]
    emotion = emotion_display[pred_idx]
    
    return f"**I detect that you are feeling {emotion}**"

# Gradio Interface
demo = gr.Interface(
    fn=predict_emotion,
    inputs=gr.Textbox(
        label="Write your message / tweet",
        lines=4,
        placeholder="I'm feeling really down today..."
    ),
    outputs=gr.Markdown(),
    title="English Tweet Emotion Classifier",
    description="TF-IDF + Logistic Regression • 6 Emotions (sadness, joy, love, anger, fear, surprise)",
    examples=[
        ["my grandfather died yesterday"],
        ["just had the best day with friends!"],
        ["this headache is killing me ugh"],
        ["finally weekend yayyyy"],
        ["i'm so scared of what might happen"],
        ["i love you so much ❤️"],
    ],
    allow_flagging="never",
    theme=gr.themes.Soft()
)

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