# app.py
import gradio as gr
from transformers import pipeline
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch
import spacy

# Load the zero-shot classification model
model_name = "MoritzLaurer/DeBERTa-v3-large-mnli-fever-anli-ling-wanli"
classifier = pipeline("zero-shot-classification", model=model_name)

# Load spaCy for brand name extraction
nlp = spacy.load("en_core_web_sm")

def extract_brand_names(text):
    doc = nlp(text)
    # Look for organization names and proper nouns that might be brands
    potential_brands = []
    for ent in doc.ents:
        if ent.label_ in ["ORG", "PRODUCT"]:
            potential_brands.append((ent.text, 0.9))  # High confidence for named entities
    
    # Also check for proper nouns that might be brands
    for token in doc:
        if token.pos_ == "PROPN" and token.text not in [brand[0] for brand in potential_brands]:
            potential_brands.append((token.text, 0.7))  # Lower confidence for proper nouns
    
    return potential_brands

def classify_product(ad_text):
    if not ad_text.strip():
        return "Please enter some ad text."
    
    # Category classification
    category_hypothesis = "This is an advertisement for a product in the category of"
    candidate_categories = [
        "Software", "Electronics", "Clothing", "Food & Beverage", 
        "Healthcare", "Financial Services", "Entertainment",
        "Home & Garden", "Automotive", "Education"
    ]
    
    category_result = classifier(
        ad_text,
        candidate_labels=candidate_categories,
        hypothesis_template=category_hypothesis,
        multi_label=False
    )
    
    # Product type classification
    product_hypothesis = "This is specifically a"
    # We'll let the model infer specific product types based on the text
    product_result = classifier(
        ad_text,
        candidate_labels=[
            "software application", "mobile app", "subscription service",
            "physical product", "digital product", "professional service",
            "consumer device", "platform", "tool"
        ],
        hypothesis_template=product_hypothesis,
        multi_label=False
    )
    
    # Brand extraction
    brands = extract_brand_names(ad_text)
    
    # Format results
    results = {
        "Category": {
            "classification": category_result["labels"][0],
            "confidence": f"{category_result['scores'][0]:.2%}"
        },
        "Product Type": {
            "classification": product_result["labels"][0],
            "confidence": f"{product_result['scores'][0]:.2%}"
        },
        "Detected Brands": [
            {"brand": brand, "confidence": f"{conf:.2%}"}
            for brand, conf in brands
        ] if brands else "No specific brand detected"
    }
    
    # Format output string
    output = f"""
📊 Analysis Results:

🏷️ Category: {results['Category']['classification']}
   Confidence: {results['Category']['confidence']}

📦 Product Type: {results['Product Type']['classification']}
   Confidence: {results['Product Type']['confidence']}

🏢 Brand Detection:"""
    
    if isinstance(results["Detected Brands"], list):
        for brand_info in results["Detected Brands"]:
            output += f"\n   • {brand_info['brand']} (Confidence: {brand_info['confidence']})"
    else:
        output += f"\n   {results['Detected Brands']}"
    
    return output

# Create Gradio interface
iface = gr.Interface(
    fn=classify_product,
    inputs=gr.Textbox(
        lines=5,
        placeholder="Paste your ad text here (max 100 words)...",
        label="Advertisement Text"
    ),
    outputs=gr.Textbox(label="Analysis Results"),
    title="AI Powered Product Identifier from Ad Text",
    description="Paste your marketing ad text to identify the product category, type, and brand mentions.",
    examples=[
        ["Experience seamless productivity with our new CloudWork Pro subscription. This AI-powered workspace solution helps remote teams collaborate better with smart document sharing, real-time editing, and integrated chat features. Starting at $29/month."],
        ["Introducing the new iPhone 15 Pro with revolutionary A17 Pro chip. Capture stunning photos with our advanced 48MP camera system. Available in titanium finish with all-day battery life. Pre-order now at Apple stores nationwide."],
    ],
    theme=gr.themes.Soft()
)