import streamlit as st
from transformers import pipeline, AutoModel, AutoTokenizer
from diffusers import StableDiffusionPipeline, DiffusionPipeline, DDIMScheduler
from datasets import load_dataset, Dataset
from peft import PeftModel, PeftConfig, get_peft_model
from accelerate import Accelerator
from optimum.onnxruntime import ORTModelForSequenceClassification
import torch
from PIL import Image
import pandas as pd

# Cache resource-intensive models for performance
@st.cache_resource
def load_diffuser():
    return StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")

# Sidebar navigation
st.sidebar.title("🤗 Hugging Face Libraries Demo")
library = st.sidebar.selectbox("Choose a library", ["Transformers", "Diffusers", "Datasets", "PEFT", "Accelerate", "Optimum"])

# Transformers Section
if library == "Transformers":
    st.header("📝 Transformers: Text Processing Powerhouse")
    st.write("Analyze or transform text with pre-trained NLP models. Try it out below! 👇")
    
    # pipeline: Easy inference for NLP tasks
    # 🎭 - Sentiment analysis with pipeline. Ex: 'I love this!' → Positive | Emotions in a snap! 😊
    # ✍️ - Text generation with pipeline. Ex: 'Once upon a time' → Story | Spin tales fast! 📖
    # ❓ - Question answering with pipeline. Ex: 'Who won?' → Answer | Quiz master ready! 🧠
    task = st.selectbox("Pick a Task", ["Sentiment Analysis 😊😢", "Text Generation ✍️", "Fill Mask 🤖"])
    text_input = st.text_area("Enter text here ✍️", "", placeholder="Type your text...")

    if st.button("Process Text 🚀", disabled=not text_input):
        with st.spinner("Processing... ⏳"):
            if "Sentiment" in task:
                classifier = pipeline("sentiment-analysis")
                result = classifier(text_input)
                st.success(f"**Sentiment Result:** {result[0]['label']} (Confidence: {result[0]['score']:.2f}) ✅")
                st.write("🎭 - Sentiment analysis with pipeline. Ex: 'I love this!' → Positive | Emotions in a snap! 😊")
            elif "Text Generation" in task:
                generator = pipeline("text-generation")
                result = generator(text_input, max_length=50)[0]['generated_text']
                st.success(f"**Generated Text:** {result} ✅")
                st.write("✍️ - Text generation with pipeline. Ex: 'Once upon a time' → Story | Spin tales fast! 📖")
            elif "Fill Mask" in task:
                fill_mask = pipeline("fill-mask")
                results = fill_mask(text_input.replace("[MASK]", fill_mask.tokenizer.mask_token))
                for res in results[:3]:
                    st.write(f"Prediction: {res['sequence']} (Score: {res['score']:.2f})")
                # AutoModel & AutoTokenizer: Load models and tokenizers
                # 🤖 - Mask filling with AutoModel. Ex: 'The [MASK] is blue' → sky | Guess master! ✨
                # 📝 - Tokenizing with AutoTokenizer. Ex: 'Hello world' → tokens | Words to numbers! 🔢
                # ⚙️ - Custom inference with AutoModel. Ex: 'BERT for sentiment' → logits | Model magic! 🪄
                st.write("🤖 - Mask filling with pipeline. Ex: 'The [MASK] is blue' → sky | Guess master! ✨")

# Diffusers Section
elif library == "Diffusers":
    st.header("🖼️ Diffusers: Image Generation Magic")
    st.write("Turn text into stunning images with diffusion models! Let's create something cool. 🎨")
    
    # StableDiffusionPipeline: Text-to-image generation
    # 🖼️ - Text-to-image with StableDiffusionPipeline. Ex: 'Cat in space' → Image | Art from words! 🚀
    # 🎨 - Style transfer with StableDiffusionPipeline. Ex: 'Van Gogh cat' → Image | Paint like pros! 🖌️
    # 🌟 - Creative prompts with StableDiffusionPipeline. Ex: 'Dream city' → Image | Imagine it, see it! ✨
    prompt = st.text_input("Enter a Creative Prompt 🎨", "", placeholder="e.g., 'A cat in space'")
    if st.button("Generate Image 🌟", disabled=not prompt):
        with st.spinner("Generating image... ⏳"):
            pipe = load_diffuser()
            image = pipe(prompt).images[0]
            st.image(image, caption="Generated Image", use_column_width=True)
            st.write("🖼️ - Text-to-image with StableDiffusionPipeline. Ex: 'Cat in space' → Image | Art from words! 🚀")
            # DiffusionPipeline & DDIMScheduler: Advanced diffusion control
            # 🧩 - Inpainting with DiffusionPipeline. Ex: 'Fix broken image' → Restored | Puzzle solver! 🛠️
            # 🔧 - Scheduling with DDIMScheduler. Ex: 'Fast diffusion' → Image | Speedy art! ⚡

# Datasets Section
elif library == "Datasets":
    st.header("📊 Datasets: Ready-to-Use Data")
    st.write("Explore datasets for training or analysis. Pick one and see a sample! 📈")
    
    # load_dataset: Access public datasets
    # 📚 - Load reviews with load_dataset. Ex: 'imdb' → Reviews | Movie buffs rejoice! 🎬
    # ❓ - Load QA with load_dataset. Ex: 'squad' → Q&A | Trivia time! 🧠
    # 🗣️ - Load audio with load_dataset. Ex: 'common_voice' → Audio | Hear the data! 🎙️
    dataset_name = st.selectbox("Choose a Dataset", ["imdb 🎬", "squad ❓"])
    if st.button("Load Dataset 📥"):
        with st.spinner("Fetching dataset... ⏳"):
            dataset = load_dataset(dataset_name.split()[0], split="train[:5]")
            st.success(f"**Dataset Loaded:** {dataset_name} (showing first 5 samples) ✅")
            st.write(dataset)
            st.write("📚 - Load reviews with load_dataset. Ex: 'imdb' → Reviews | Movie buffs rejoice! 🎬")

    # Dataset.from_pandas: Convert DataFrames to datasets
    # 📊 - CSV to dataset with Dataset.from_pandas. Ex: 'Tweets CSV' → Dataset | Your data shines! 🐦
    # 📈 - Analyze data with Dataset.from_pandas. Ex: 'Sales data' → Dataset | Numbers talk! 💰
    # 🔄 - Preprocess with Dataset.from_pandas. Ex: 'Raw text' → Dataset | Clean it up! 🧹
    st.subheader("Create Dataset from CSV")
    uploaded_file = st.file_uploader("Upload a CSV file", type="csv")
    if uploaded_file is not None:
        df = pd.read_csv(uploaded_file)
        dataset = Dataset.from_pandas(df)
        st.write(dataset)
        st.write("📊 - CSV to dataset with Dataset.from_pandas. Ex: 'Tweets CSV' → Dataset | Your data shines! 🐦")

# PEFT Section (Simplified)
elif library == "PEFT":
    st.header("⚙️ PEFT: Efficient Fine-Tuning")
    st.write("Learn about parameter-efficient fine-tuning (simplified demo).")
    
    # PeftModel: Load fine-tuned models
    # ⚙️ - Inference with PeftModel. Ex: 'Adapted BERT' → Output | Slim yet mighty! 💪
    # 🤖 - Custom tasks with PeftModel. Ex: 'NER tuning' → Tags | Precision tweak! 🎯
    # 📉 - Low resource with PeftModel. Ex: 'Small GPU' → Model | Efficiency wins! 🏆
    # PeftConfig: Configure PEFT settings
    # 🛠️ - LoRA setup with PeftConfig. Ex: 'Rank 8' → Config | Tune light! 🌟
    # 🔧 - Adapter config with PeftConfig. Ex: 'Task-specific' → Config | Fit like a glove! 🧤
    # ⚡ - Fast tuning with PeftConfig. Ex: 'Quick fine-tune' → Config | Speedy prep! 🏃
    # get_peft_model: Wrap models with PEFT
    # 🔗 - Enhance with get_peft_model. Ex: 'BERT + LoRA' → Model | Power boost! 🚀
    # 📝 - Task adapt with get_peft_model. Ex: 'Sentiment' → Model | Tailored fit! ✂️
    # 🌐 - Multi-task with get_peft_model. Ex: 'QA + NER' → Model | Versatility rules! 🌍
    text = st.text_area("Enter Text to Classify ✍️", "", placeholder="Type something...")
    if st.button("Classify 🚀", disabled=not text):
        with st.spinner("Processing... ⏳"):
            st.success("**Result:** Placeholder (PEFT reduces parameters efficiently!) ✅")
            st.write("⚙️ - Inference with PeftModel. Ex: 'Adapted BERT' → Output | Slim yet mighty! 💪")

# Accelerate Section (Simplified)
elif library == "Accelerate":
    st.header("🚀 Accelerate: Device Optimization")
    st.write("Optimize inference across devices with Accelerate.")
    
    # Accelerator: Manage device placement
    # 🚀 - GPU use with Accelerator. Ex: 'Auto GPU' → Device | Speed unleashed! ⚡
    # 🔗 - Multi-device with Accelerator. Ex: '2 GPUs' → Setup | Teamwork rocks! 🤝
    # 🛠️ - Easy setup with Accelerator. Ex: 'Model prep' → Ready | Smooth start! 🌈
    text = st.text_area("Enter Text for Sentiment Analysis ✍️", "", placeholder="Type something...")
    if st.button("Analyze 🚀", disabled=not text):
        with st.spinner("Analyzing... ⏳"):
            accelerator = Accelerator()
            classifier = pipeline("sentiment-analysis")
            result = classifier(text)
            st.success(f"**Result:** {result[0]['label']} (Confidence: {result[0]['score']:.2f}) ✅")
            st.write("🚀 - GPU use with Accelerator. Ex: 'Auto GPU' → Device | Speed unleashed! ⚡")

# Optimum Section (Simplified)
elif library == "Optimum":
    st.header("⚡ Optimum: Hardware Acceleration")
    st.write("Speed up inference with optimized models (e.g., ONNX).")
    
    # ORTModelForSequenceClassification: ONNX-optimized classification
    # ⚡ - Fast classify with ORTModelForSequenceClassification. Ex: 'Text' → Label | Speed king! 🏎️
    # 🖥️ - CPU boost with ORTModelForSequenceClassification. Ex: 'Sentiment' → Result | No GPU? No prob! 💪
    # 📊 - Efficient with ORTModelForSequenceClassification. Ex: 'Batch text' → Labels | Quick wins! 🏆
    text = st.text_area("Enter Text to Classify ✍️", "", placeholder="Type something...")
    if st.button("Classify 🚀", disabled=not text):
        with st.spinner("Processing... ⏳"):
            st.success("**Result:** Placeholder (Optimum boosts speed!) ✅")
            st.write("⚡ - Fast classify with ORTModelForSequenceClassification. Ex: 'Text' → Label | Speed king! 🏎️")