import gradio as gr
from transformers import pipeline
import numpy as np

# Initialize the pipeline
pipe = pipeline("text-classification", model="AbrorBalxiyev/text-classification", return_all_scores=True)

label_mapping = {
    0: 'Avto', 1: 'Biznes', 2: 'Iqtisodiyot', 3: 'Kino', 4: 'Kitob',
    5: 'Koinot', 6: 'Madaniyat', 7: 'Ob-havo', 8: 'Sayohat', 9: 'Sport', 10: 'Texnologiya'
}

def get_html_for_results(results):
    # Sort results by score in descending order
    sorted_results = sorted(results, key=lambda x: x['score'], reverse=True)
    
    html = """
    <style>
        .result-container {
            font-family: Arial, sans-serif;
            max-width: 600px;
            margin: 20px auto;
        }
        .category-row {
            margin: 10px 0;
        }
        .category-name {
            display: inline-block;
            width: 120px;
            font-size: 14px;
            color: #333;
        }
        .progress-bar {
            display: inline-block;
            width: calc(100% - 200px);
            height: 20px;
            background-color: #f0f0f0;
            border-radius: 10px;
            overflow: hidden;
            margin-right: 10px;
        }
        .progress {
            height: 100%;
            background-color: #ff6b33;
            border-radius: 10px;
            transition: width 0.5s ease-in-out;
        }
        .percentage {
            display: inline-block;
            width: 50px;
            text-align: right;
            color: #666;
        }
    </style>
    <div class="result-container">
    """
    
    for item in sorted_results:
        percentage = item['score'] * 100
        html += f"""
        <div class="category-row">
            <span class="category-name">{item['label']}</span>
            <div class="progress-bar">
                <div class="progress" style="width: {percentage}%;"></div>
            </div>
            <span class="percentage">{percentage:.0f}%</span>
        </div>
        """
    
    html += "</div>"
    return html

def classify_text(text):
    if not text.strip():
        return "Please enter some text to classify."
    
    # Get predictions
    pred = pipe(text)
    
    # Decode predictions
    decoded_data = [
        {"label": label_mapping[int(item["label"].split("_")[1])], 
         "score": item["score"]} for item in pred[0]
    ]
    
    return get_html_for_results(decoded_data)

# Create Gradio interface
iface = gr.Interface(
    fn=classify_text,
    inputs=[
        gr.Textbox(
            placeholder="Enter text to classify...",
            label=None,
            lines=3
        )
    ],
    outputs=gr.HTML(),
    title="Text Category Classification",
    css="""
        .gradio-container {
            font-family: Arial, sans-serif;
        }
        .gradio-interface {
            max-width: 800px !important;
        }
        #component-0 {
            border-radius: 8px;
            border: 1px solid #ddd;
        }
        .submit-button {
            background-color: #ff6b33 !important;
        }
        .clear-button {
            background-color: #f0f0f0 !important;
            color: #333 !important;
        }
    """,
    examples=[
        ["Messi jahon chempioni bo'ldi"],
        ["Yangi iPhone 15 Pro Max sotuvga chiqdi"],
        ["Kitob o'qish foydali"],
        ["Toshkentda ob-havo issiq"]
    ]
)

# Launch the interface
iface.launch(share=True)