import torch
import torch.nn.functional as F
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from bertopic import BERTopic
from sentence_transformers import SentenceTransformer
import os

# --- 1. CONFIGURATION ---
# Tránh xung đột luồng khi chạy tokenizer
os.environ["TOKENIZERS_PARALLELISM"] = "false"

ROBERTA_PATH = "roberta_base_checkpoint"
BERTOPIC_PATH = "bertopic_checkpoint"

# --- 2. INITIALIZATION (Runs once at startup) ---
print("--- System Initializing ---")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device detected: {device.type.upper()}")

# Load RoBERTa
tokenizer = AutoTokenizer.from_pretrained(ROBERTA_PATH)
roberta_model = AutoModelForSequenceClassification.from_pretrained(ROBERTA_PATH).to(device)
roberta_model.eval()

# Load BERTopic
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
topic_model = BERTopic.load(BERTOPIC_PATH, embedding_model=embedding_model)

print("--- All models loaded successfully! ---")

# --- 3. CORE LOGIC ---
def analyze_review(text):
    if not text.strip():
        return "Please enter some text", None, None, ""

    # --- Sentiment Analysis (RoBERTa) ---
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=256).to(device)
    with torch.no_grad():
        outputs = roberta_model(**inputs)
    
    probs = F.softmax(outputs.logits, dim=1)
    pred_idx = torch.argmax(probs, dim=1).item()
    
    # Matching common sentiment output formats
    sentiment_labels = ['Negative 😡', 'Neutral 😐', 'Positive 😊']
    sentiment = sentiment_labels[pred_idx]
    confidence = probs[0][pred_idx].item()

    # --- Topic Modeling (BERTopic) ---
    new_topics, _ = topic_model.transform([text])
    topic_id = int(new_topics[0])
    
    if topic_id != -1:
        keywords_data = topic_model.get_topic(topic_id)
        # Extracting top 5 keywords
        keywords_list = ", ".join([word for word, score in keywords_data[:5]])
    else:
        keywords_list = "Outlier (No specific topic detected)"

    return sentiment, f"{confidence:.2%}", topic_id, keywords_list

# --- 4. GRADIO INTERFACE (WEB & API) ---
with gr.Blocks(title="Food Review Analysis Tool") as demo:
    gr.Markdown("# 📊 Food Review Analysis System")
    gr.Markdown("An integrated pipeline using **RoBERTa** for Sentiment and **BERTopic** for Topic Extraction.")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="Review Content", 
                placeholder="e.g., The pizza was delicious but the delivery took two hours!",
                lines=5
            )
            btn = gr.Button("Analyze Now", variant="primary")
            
        with gr.Column():
            out_sentiment = gr.Label(label="Sentiment Prediction")
            out_confidence = gr.Textbox(label="Confidence Score")
            with gr.Row():
                out_topic_id = gr.Number(label="Topic ID")
                out_keywords = gr.Textbox(label="Main Keywords")

    # Event binding
    btn.click(
        fn=analyze_review, 
        inputs=input_text, 
        outputs=[out_sentiment, out_confidence, out_topic_id, out_keywords]
    )
    
    # Examples for quick testing
    gr.Examples(
        examples=[
            ["The food was amazing, very delicious!"],
            ["Service was very slow and the waiter was rude."],
            ["While the truffle-infused risotto was a masterpiece, the deafening noise made it hard to enjoy."]
        ],
        inputs=input_text
    )

# Launch the app
if __name__ == "__main__":
    demo.launch()