import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch.nn.functional as F

# Lightweight model already fine-tuned for sentiment (91%+ accuracy)
# Much smaller than bert-base-uncased → fits easily in free CPU Spaces
MODEL_NAME = "distilbert/distilbert-base-uncased-finetuned-sst-2-english"

# Load model efficiently
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

model = AutoModelForSequenceClassification.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float32,   # Safe for CPU
    low_cpu_mem_usage=True
)

model.eval()

# Safe device handling
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

def predict_sentiment(text: str):
    if not text or not text.strip():
        return "Please enter some text", 0.0, "⚠️"

    # Tokenize
    inputs = tokenizer(
        text, 
        return_tensors="pt", 
        truncation=True, 
        padding=True, 
        max_length=512
    ).to(device)

    # Inference with no gradient
    with torch.no_grad():
        outputs = model(**inputs)
        probs = F.softmax(outputs.logits, dim=-1)

    # Get prediction and confidence
    pred = torch.argmax(probs, dim=-1).item()
    confidence = probs[0][pred].item() * 100

    if pred == 1:  # 1 = positive in this model
        sentiment = "Positive 😊"
        emoji = "🟢"
    else:
        sentiment = "Negative 😞"
        emoji = "🔴"

    return sentiment, round(confidence, 2), emoji

# ====================== Gradio UI ======================
with gr.Blocks(theme=gr.themes.Soft(), title="BERT Sentiment Analysis") as demo:
    gr.Markdown("# 🎬 Transformer Sentiment Analysis")
    gr.Markdown("**DistilBERT** (lightweight BERT) for fast movie review / text sentiment prediction")

    with gr.Row():
        text_input = gr.Textbox(
            label="Enter your text (movie review, comment, tweet, etc.)",
            placeholder="Type or paste here...",
            lines=5,
        )

    analyze_btn = gr.Button("🔍 Analyze Sentiment", variant="primary", size="large")

    with gr.Row():
        sentiment_output = gr.Textbox(label="Sentiment Result", interactive=False)
        confidence_output = gr.Number(label="Confidence (%)", interactive=False)
        indicator_output = gr.Textbox(label="Indicator", interactive=False)

    # Nice examples
    gr.Examples(
        examples=[
            ["This movie was absolutely fantastic! The acting and story were brilliant."],
            ["I really disliked the plot. It was boring and predictable."],
            ["The visuals were great but the dialogue felt terrible."],
            ["One of the best films I have watched this year. Highly recommended!"],
            ["Complete waste of time. Do not watch this movie."],
        ],
        inputs=text_input,
        outputs=[sentiment_output, confidence_output, indicator_output],
        fn=predict_sentiment,
        cache_examples=False,
    )

    # Button click
    analyze_btn.click(
        fn=predict_sentiment,
        inputs=text_input,
        outputs=[sentiment_output, confidence_output, indicator_output],
    )

# Launch (Important: share=False on HF Spaces)
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True,
        debug=False
    )