# app.py — WILFRID Emotion Oracle — WORKS ON HUGGING FACE SPACES 100%

import os
from huggingface_hub import login
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import torch
import gradio as gr
from emoji import emojize

# Authenticate with your secret token
HF_TOKEN = os.getenv("HF_TOKEN")
if HF_TOKEN:
    login(token=HF_TOKEN)
    print("Logged in to Hugging Face")
else:
    raise RuntimeError("HF_TOKEN secret not found! Add it in Settings → Secrets")

print("Loading Llama-3-8B-Instruct + WILFRID adapter...")
# Load base model — no offloading, full GPU
base_model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B-Instruct",
    torch_dtype=torch.bfloat16,
    device_map="auto",
    low_cpu_mem_usage=True,
    offload_folder=None,
    offload_state_dict=False,
)

# Load your LoRA adapter
model = PeftModel.from_pretrained(base_model, "Wilfrid28/llama3-meld-wilfrid")
tokenizer = AutoTokenizer.from_pretrained("Wilfrid28/llama3-meld-wilfrid", use_fast=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model.eval()
print("WILFRID is ready — Emotion Oracle activated!")

# Emotion token IDs
EMOTION_IDS = {
    e: tokenizer(f" {e}", return_tensors="pt")["input_ids"][0,1].item()
    for e in ["neutral","joy","surprise","anger","sadness","disgust","fear"]
}

EMOJI_MAP = {
    "joy": ":beaming_face_with_smiling_eyes:", "anger": ":angry_face:",
    "sadness": ":loudly_crying_face:", "surprise": ":face_screaming_in_fear:",
    "fear": ":fearful_face:", "disgust": ":nauseated_face:", "neutral": ":neutral_face:"
}

def predict(message, history):
    if not message.strip():
        return history

    conv = "\n".join([f"You: {u}" for u, _ in history]) + f"\nYou: {message}\n"
    prompt = f"Conversation:\n{conv}Current turn:\nYou: {message}\nEmotion (one word):"

    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        logits = model(**inputs).logits[0, -1]
        probs = torch.softmax(torch.tensor([logits[EMOTION_IDS[e]] for e in EMOTION_IDS]), dim=0)

    top3 = torch.topk(probs, 3)
    emotions = ["neutral","joy","surprise","anger","sadness","disgust","fear"]
    top_emotion = emotions[top3.indices[0]]
    confidence = top3.values[0].item() * 100
    emoji = emojize(EMOJI_MAP[top_emotion])

    reply = f"""{emoji} **{top_emotion.upper()}** — {confidence:.1f}%

**Top 3**
• {emotions[top3.indices[0]].capitalize()} — {top3.values[0]*100:.1f}%
• {emotions[top3.indices[1]].capitalize()} — {top3.values[1]*100:.1f}%
• {emotions[top3.indices[2]].capitalize()} — {top3.values[2]*100:.1f}%"""

    history.append((message, reply))
    return history

# UI — CLEAN AND WORKING
with gr.Blocks() as demo:
    gr.Markdown("# WILFRID — 0.6607 F1 Emotion Oracle")
    chatbot = gr.Chatbot(height=600)
    msg = gr.Textbox(placeholder="Type your message...", container=False, scale=7)
    send = gr.Button("Send", variant="primary")

    msg.submit(predict, [msg, chatbot], chatbot)
    send.click(predict, [msg, chatbot], chatbot)

print("WILFRID is launching...")
demo.launch()