# Install required packages
#!pip install transformers gradio torch sentencepiece

import gradio as gr
from transformers import pipeline

print("🚀 Setting up AI pipelines...")

# Initialize all pipelines with specific models for better performance
pipelines = {
    "sentiment": pipeline("sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english"),
    "text_generation": pipeline("text-generation", model="gpt2", max_length=100),
    "qa": pipeline("question-answering", model="distilbert-base-cased-distilled-squad"),
    "translation": pipeline("translation_en_to_fr", model="t5-small"),
    "summarization": pipeline("summarization", model="facebook/bart-large-cnn"),
    "ner": pipeline("ner", aggregation_strategy="simple", model="dbmdz/bert-large-cased-finetuned-conll03-english"),
    "zero_shot": pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
}

print("✅ All pipelines loaded successfully!")

def ai_playground(task, text, context, labels):
    try:
        if task == "sentiment":
            result = pipelines["sentiment"](text)[0]
            emotion = "😊" if result['label'] == 'POSITIVE' else '😞'
            return f"{emotion} {result['label']} ({result['score']:.2%} confidence)"

        elif task == "text_generation":
            result = pipelines["text_generation"](text, max_length=100, do_sample=True)[0]['generated_text']
            return f"📝 {result}"

        elif task == "qa":
            if not context.strip():
                return "❌ Please provide context for question answering"
            result = pipelines["qa"](question=text, context=context)
            return f"✅ Answer: {result['answer']}\n🎯 Confidence: {result['score']:.2%}"

        elif task == "translation":
            result = pipelines["translation"](text)[0]['translation_text']
            return f"🇫🇷 {result}"

        elif task == "summarization":
            if len(text) < 50:
                return "❌ Text is too short for summarization. Please provide longer text."
            result = pipelines["summarization"](text, max_length=80, min_length=30, do_sample=False)[0]['summary_text']
            return f"📌 Summary: {result}"

        elif task == "ner":
            entities = pipelines["ner"](text)
            if not entities:
                return "🔍 No named entities found"
            result = "\n".join([f"• {entity['word']} → {entity['entity_group']} ({entity['score']:.2%})" for entity in entities[:10]])
            return f"🔍 Found Entities:\n{result}"

        elif task == "zero_shot":
            if not labels.strip():
                return "❌ Please provide labels for zero-shot classification"
            label_list = [label.strip() for label in labels.split(",") if label.strip()]
            result = pipelines["zero_shot"](text, label_list, multi_label=False)
            formatted_results = "\n".join([f"• {label}: {score:.2%}" for label, score in zip(result['labels'][:5], result['scores'][:5])])
            return f"🏷️ Classification:\n{formatted_results}"

        else:
            return "Please provide required inputs"

    except Exception as e:
        return f"❌ Error: {str(e)}"

# Create the interface with better layout
with gr.Blocks(theme=gr.themes.Soft(), title="🎪 AI Playground") as demo:
    gr.Markdown("# 🎪 AI Playground - Try Everything!")
    gr.Markdown("Perfect for beginners! Choose a task and see AI in action. Built for Google Colab! 🚀")

    with gr.Row():
        with gr.Column(scale=1):
            task_dropdown = gr.Dropdown(
                choices=["sentiment", "text_generation", "qa", "translation", "summarization", "ner", "zero_shot"],
                label="🎯 Choose AI Task",
                value="sentiment",
                info="Select what you want to do"
            )

            input_text = gr.Textbox(
                lines=3,
                placeholder="Enter your text here...",
                label="📝 Input Text",
                value="I love learning about artificial intelligence!"
            )

            context_text = gr.Textbox(
                lines=3,
                placeholder="For QA: Enter context here...\nExample: Hugging Face provides AI tools and models. The company was founded in 2016.",
                label="📚 Context (for QA only)",
                visible=False
            )

            labels_text = gr.Textbox(
                lines=2,
                placeholder="For Zero-shot: Enter comma-separated labels\nExample: positive, negative, neutral",
                label="🏷️ Labels (for Zero-shot only)",
                visible=False
            )

            submit_btn = gr.Button("✨ Run AI Magic!", variant="primary")

        with gr.Column(scale=2):
            output_text = gr.Textbox(
                lines=8,
                label="🎉 AI Output",
                interactive=False,
                show_copy_button=True
            )

    # Examples for each task
    gr.Markdown("## 💡 Try These Examples:")

    examples = [
        ["sentiment", "I'm so excited to learn about AI!", "", ""],
        ["text_generation", "Once upon a time in a magical kingdom,", "", ""],
        ["translation", "Hello, how are you today?", "", ""],
        ["summarization", "Artificial intelligence is transforming many industries. Machine learning allows computers to learn from data. Deep learning uses neural networks for complex tasks. AI is used in healthcare, finance, and transportation.", "", ""],
        ["ner", "Apple was founded by Steve Jobs in Cupertino, California in 1976.", "", ""],
        ["zero_shot", "The new smartphone has amazing battery life and camera quality.", "", "technology, review, complaint, inquiry"],
        ["qa", "What does Hugging Face provide?", "Hugging Face is a company that provides AI tools and models for natural language processing. Their library makes it easy to use state-of-the-art models.", ""]
    ]

    gr.Examples(
        examples=examples,
        inputs=[task_dropdown, input_text, context_text, labels_text],
        outputs=output_text,
        fn=ai_playground,
        cache_examples=False,
        label="Click any example below to try it!"
    )

    # Show/hide context and labels based on task selection
    def toggle_inputs(task):
        context_visible = (task == "qa")
        labels_visible = (task == "zero_shot")
        return [
            gr.Textbox(visible=context_visible),
            gr.Textbox(visible=labels_visible)
        ]

    task_dropdown.change(
        fn=toggle_inputs,
        inputs=task_dropdown,
        outputs=[context_text, labels_text]
    )

    # Connect the submit button
    submit_btn.click(
        fn=ai_playground,
        inputs=[task_dropdown, input_text, context_text, labels_text],
        outputs=output_text
    )

    gr.Markdown("---")
    gr.Markdown("### 🎓 Learning Tips:")
    gr.Markdown("""
    - **Sentiment Analysis**: Detects positive/negative emotions in text
    - **Text Generation**: Creates new text based on your prompt
    - **Translation**: Translates English to French
    - **Summarization**: Shortens long text while keeping key points
    - **Named Entity Recognition**: Finds people, places, organizations in text
    - **Zero-Shot Classification**: Classifies text into custom categories
    - **Question Answering**: Answers questions based on provided context
    """)

print("🎉 AI Playground is ready! Launching interface...")
demo.launch(share=True, debug=True)