import torch
import gradio as gr
from transformers import (
    pipeline,
    BartForConditionalGeneration,
    BartTokenizer
)

# ── Device ────────────────────────────────────
DEVICE     = 0 if torch.cuda.is_available() else -1
DEVICE_STR = "cuda" if torch.cuda.is_available() else "cpu"

# ── Categories & Aspects ──────────────────────
CATEGORIES = [
    "Electronics",
    "Clothing and Fashion",
    "Books and Literature",
    "Food and Grocery",
    "Sports and Outdoors",
    "Home and Kitchen",
    "Beauty and Personal Care",
    "Toys and Games"
]

ASPECTS = [
    "price and value",
    "quality and durability",
    "delivery and shipping",
    "customer service",
    "ease of use",
    "design and appearance",
    "performance and speed",
    "battery life"
]

# ── Load all models from HuggingFace Hub ──────
print("Loading sentiment model...")
sentiment_pipe = pipeline(
    "text-classification",
    model  = "Ved2001/pranalyzer",
    device = DEVICE
)

print("Loading category classifier...")
category_pipe = pipeline(
    "zero-shot-classification",
    model  = "facebook/bart-large-mnli",
    device = DEVICE
)

print("Loading aspect analyzer...")
aspect_pipe = pipeline(
    "zero-shot-classification",
    model  = "cross-encoder/nli-roberta-base",
    device = DEVICE
)

print("Loading summarization model...")
bart_tokenizer = BartTokenizer.from_pretrained(
    "facebook/bart-large-xsum")
bart_model = BartForConditionalGeneration\
    .from_pretrained("facebook/bart-large-xsum")\
    .to(DEVICE_STR)

print("All models loaded!")


# ── Inference functions ───────────────────────
def analyze_sentiment(text):
    result = sentiment_pipe(text[:512])[0]
    return {'label': result['label'], 'score': round(result['score'], 4)}

def classify_category(text):
    result = category_pipe(
        text[:512],
        candidate_labels=CATEGORIES,
        multi_label=False
    )
    return {'category': result['labels'][0], 'score': round(result['scores'][0], 4)}

def analyze_aspects(text):
    result = aspect_pipe(
        text[:512],
        candidate_labels=ASPECTS,
        multi_label=True
    )
    return [
        (label, round(score, 4))
        for label, score in zip(result['labels'], result['scores'])
        if score > 0.3
    ][:3]

def summarize_review(text):
    if len(text.split()) < 30:
        return text
    inputs = bart_tokenizer(
        text[:512],
        return_tensors="pt",
        truncation=True,
        max_length=512
    ).to(DEVICE_STR)
    summary_ids = bart_model.generate(
        inputs["input_ids"],
        max_new_tokens=80,
        min_length=15,
        length_penalty=2.0,
        num_beams=4,
        early_stopping=True
    )
    return bart_tokenizer.decode(
        summary_ids[0], skip_special_tokens=True)


# ── Gradio function ───────────────────────────
def run_analysis(review_text):
    if not review_text.strip():
        return ("Please enter a review!",) * 4

    sentiment = analyze_sentiment(review_text)
    category  = classify_category(review_text)
    aspects   = analyze_aspects(review_text)
    summary   = summarize_review(review_text)

    # Format sentiment
    emoji = "😊" if sentiment['label'] == "POSITIVE" else "😞"
    sentiment_out = (
        f"{emoji} {sentiment['label']}\n"
        f"Confidence: {sentiment['score']*100:.1f}%"
    )

    # Format category
    category_out = (
        f"📦 {category['category']}\n"
        f"Confidence: {category['score']*100:.1f}%"
    )

    # Format aspects
    aspects_out = "🔍 Aspects Mentioned:\n\n"
    for aspect, score in aspects:
        bar   = "█" * int(score * 10)
        empty = "░" * (10 - int(score * 10))
        aspects_out += f"• {aspect:<25} {score:.2f}  {bar}{empty}\n"
    if not aspects:
        aspects_out += "No strong aspects detected."

    # Format summary
    summary_out = f"📝 {summary}"

    return sentiment_out, category_out, aspects_out, summary_out


# ── Gradio UI ─────────────────────────────────
with gr.Blocks(title="pranalyzer") as demo:

    gr.Markdown("""
    # 🛍️ pranalyzer
    ### Product Review Analyzer
    Paste any Amazon product review and get instant analysis
    using 4 NLP models running in parallel.
    ---
    """)

    review_input = gr.Textbox(
        label       = "📝 Paste your product review here",
        placeholder = "e.g. This laptop is amazing! Battery lasts all day...",
        lines       = 6
    )

    analyze_btn = gr.Button(
        "🔍 Analyze Review",
        variant = "primary",
        size    = "lg"
    )

    gr.Markdown("### 📊 Analysis Results")

    with gr.Row():
        sentiment_out = gr.Textbox(label="😊 Sentiment", lines=3)
        category_out  = gr.Textbox(label="📦 Category",  lines=3)

    with gr.Row():
        aspects_out = gr.Textbox(label="🔍 Aspects", lines=6)
        summary_out = gr.Textbox(label="📝 Summary", lines=6)

    gr.Examples(
        examples=[
            ["This laptop is absolutely incredible! Battery lasts all day, easily 10-12 hours of real work. The display is crisp and bright. Performance is blazing fast. Highly recommended!"],
            ["Complete waste of money. Stopped working after a week. Customer service was useless and refused a refund. Avoid at all costs."],
            ["Ordered these running shoes for marathon training. Delivery was super fast. Cushioning is excellent. Only downside is sizing runs small, order a size up."],
            ["This cookbook is a disappointment. Half the recipes have missing ingredients. Very misleading. Wasted expensive ingredients trying three different recipes."]
        ],
        inputs=review_input
    )

    gr.Markdown("""
    ---
    **Models:** `DistilBERT` → Sentiment | `BART-MNLI` → Category | `RoBERTa` → Aspects | `BART-XSUM` → Summary  
    Built by [Vedant Nagarkar](https://huggingface.co/Ved2001) • 
    [GitHub](https://github.com/Vedant-Nagarkar/product-review-analyzer)
    """)

    analyze_btn.click(
        fn      = run_analysis,
        inputs  = review_input,
        outputs = [sentiment_out, category_out, aspects_out, summary_out]
    )

demo.launch()