import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import time

# Load model and tokenizer
def load_model():
    tokenizer = AutoTokenizer.from_pretrained("VietAI/vit5-base-vietnews-summarization")
    model = AutoModelForSeq2SeqLM.from_pretrained("VietAI/vit5-base-vietnews-summarization")
    return tokenizer, model

tokenizer, model = load_model()

def summarize_text(text, max_length, min_length, num_beams):
    if not text.strip():
        return "⚠️ Vui lòng nhập văn bản cần tóm tắt!", "", ""
    
    start_time = time.time()
    
    try:
        # Tokenize input
        input_ids = tokenizer.encode(
            "summarize: " + text, 
            return_tensors="pt", 
            max_length=512, 
            truncation=True
        )
        
        # Generate summary
        output_ids = model.generate(
            input_ids,
            max_length=max_length,
            min_length=min_length,
            length_penalty=2.0,
            num_beams=num_beams,
            early_stopping=True
        )
        
        # Decode result
        summary = tokenizer.decode(output_ids[0], skip_special_tokens=True)
        
        end_time = time.time()
        processing_time = round(end_time - start_time, 2)
        
        # Calculate statistics
        original_words = len(text.split())
        summary_words = len(summary.split())
        compression_ratio = round((1 - summary_words / original_words) * 100, 1)
        
        stats = f"""
📊 **Thống kê:**
- Văn bản gốc: {original_words} từ
- Bản tóm tắt: {summary_words} từ  
- Tỷ lệ nén: {compression_ratio}%
- Thời gian xử lý: {processing_time}s
        """
        
        return summary, stats, "✅ Tóm tắt thành công!"
        
    except Exception as e:
        return f"❌ Có lỗi xảy ra: {str(e)}", "", "❌ Lỗi xử lý"

# Sample texts
sample_texts = {
    "Tin tức công nghệ": """Công nghệ trí tuệ nhân tạo (AI) đang phát triển với tốc độ chóng mặt và tạo ra những thay đổi sâu sắc trong nhiều lĩnh vực của đời sống. Từ y tế, giáo dục, giao thông đến tài chính, AI đang được ứng dụng rộng rãi để cải thiện hiệu quả và chất lượng dịch vụ. 

Trong lĩnh vực y tế, AI giúp chẩn đoán bệnh chính xác hơn thông qua phân tích hình ảnh y khoa và dữ liệu bệnh nhân. Các hệ thống AI có thể phát hiện ung thư sớm, dự đoán nguy cơ mắc bệnh và đề xuất phương án điều trị phù hợp.

Tuy nhiên, sự phát triển nhanh chóng này cũng đặt ra nhiều thách thức về đạo đức, bảo mật thông tin và tác động đến thị trường lao động. Việc đảm bảo AI phát triển một cách có trách nhiệm và minh bạch là ưu tiên hàng đầu.""",
    
    "Tin tức kinh tế": """Nền kinh tế Việt Nam trong quý đầu năm 2024 ghi nhận nhiều tín hiệu tích cực với tốc độ tăng trưởng GDP đạt 6.14%, cao hơn so với cùng kỳ năm trước. Xuất khẩu tiếp tục là động lực chính thúc đẩy tăng trưởng, đặc biệt trong các ngành công nghiệp chế biến chế tạo và nông sản.

Khu vực công nghiệp và xây dựng tăng 7.29%, trong đó công nghiệp chế biến, chế tạo tăng 7.99%. Khu vực dịch vụ phục hồi mạnh mẽ với mức tăng 6.93%, được thúc đẩy bởi du lịch và thương mại.

Tuy nhiên, lạm phát vẫn là thách thức cần được kiểm soát chặt chẽ để đảm bảo ổn định kinh tế vĩ mô. Chính phủ đang thực hiện các biện pháp điều hành linh hoạt để duy trì tăng trưởng bền vững.""",
    
    "Tin tức giáo dục": """Bộ Giáo dục và Đào tạo vừa công bố kế hoạch đổi mới giáo dục toàn diện, tập trung vào phát triển năng lực và phẩm chất học sinh. Chương trình giáo dục mới nhấn mạnh việc giảm tải nội dung, tăng cường thực hành và ứng dụng kiến thức vào thực tế.

Đặc biệt, việc tích hợp công nghệ thông tin vào giảng dạy được đẩy mạnh với việc trang bị thiết bị hiện đại cho các trường học. Giáo viên được đào tạo sử dụng các công cụ số để nâng cao chất lượng bài giảng.

Bên cạnh đó, chương trình giáo dục hướng nghiệp được chú trọng để học sinh có định hướng rõ ràng về tương lai. Các trường đại học cũng liên kết chặt chẽ với doanh nghiệp để tạo cơ hội thực tập và việc làm cho sinh viên."""
}

def load_sample(sample_name):
    return sample_texts.get(sample_name, "")

# Custom CSS
css = """
.gradio-container {
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
}

.main-header {
    text-align: center;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    color: white;
    padding: 2rem;
    border-radius: 15px;
    margin-bottom: 2rem;
    box-shadow: 0 8px 32px rgba(0, 0, 0, 0.1);
}

.feature-box {
    background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%);
    color: white;
    padding: 1.5rem;
    border-radius: 10px;
    margin: 1rem 0;
    box-shadow: 0 4px 15px rgba(0, 0, 0, 0.1);
}

.stats-box {
    background: #f8f9fa;
    border: 1px solid #dee2e6;
    border-radius: 8px;
    padding: 1rem;
    margin: 0.5rem 0;
}

footer {
    text-align: center;
    padding: 2rem;
    color: #6c757d;
    border-top: 1px solid #dee2e6;
    margin-top: 2rem;
}
"""

# Create Gradio interface
with gr.Blocks(css=css, title="Tóm tắt văn bản tiếng Việt", theme=gr.themes.Soft()) as app:
    
    # Header
    gr.HTML("""
    <div class="main-header">
        <h1>📝 Tóm tắt văn bản tiếng Việt với AI</h1>
        <p style="font-size: 1.2em; margin: 0.5rem 0;">Powered by VietAI T5 Model</p>
        <p style="font-size: 1em; margin: 0;">Công cụ tóm tắt văn bản tiếng Việt thông minh và chính xác</p>
    </div>
    """)
    
    with gr.Row():
        # Left column - Main content
        with gr.Column(scale=2):
            gr.Markdown("## ✍️ Nhập văn bản cần tóm tắt")
            
            # Sample text selector
            sample_dropdown = gr.Dropdown(
                choices=list(sample_texts.keys()),
                label="📚 Chọn văn bản mẫu (tùy chọn)",
                value=None
            )
            
            # Text input
            text_input = gr.Textbox(
                label="📝 Văn bản đầu vào",
                placeholder="Nhập hoặc dán văn bản tiếng Việt cần tóm tắt vào đây...",
                lines=10,
                max_lines=15
            )
            
            # Load sample text
            sample_dropdown.change(
                fn=load_sample,
                inputs=sample_dropdown,
                outputs=text_input
            )
        
        # Right column - Settings
        with gr.Column(scale=1):
            gr.Markdown("## ⚙️ Cài đặt tóm tắt")
            
            max_length = gr.Slider(
                minimum=50,
                maximum=300,
                value=150,
                step=10,
                label="📏 Độ dài tối đa"
            )
            
            min_length = gr.Slider(
                minimum=20,
                maximum=100,
                value=40,
                step=5,
                label="📐 Độ dài tối thiểu"
            )
            
            num_beams = gr.Dropdown(
                choices=[2, 4, 6, 8],
                value=4,
                label="🔍 Số beam search"
            )
            
            # Model info
            gr.HTML("""
            <div class="feature-box">
                <h4>🤖 Thông tin Model</h4>
                <ul style="margin: 0; padding-left: 1.2rem;">
                    <li><strong>Model:</strong> VietAI T5</li>
                    <li><strong>Ngôn ngữ:</strong> Tiếng Việt</li>
                    <li><strong>Loại:</strong> Seq2Seq</li>
                </ul>
            </div>
            """)
            
            # Tips
            gr.HTML("""
            <div class="stats-box">
                <h4>💡 Mẹo sử dụng</h4>
                <ul style="margin: 0.5rem 0; padding-left: 1.2rem; font-size: 0.9em;">
                    <li>Văn bản 100-1000 từ cho kết quả tốt nhất</li>
                    <li>Sử dụng tiếng Việt chuẩn</li>
                    <li>Tránh quá nhiều từ viết tắt</li>
                </ul>
            </div>
            """)
    
    # Summarize button
    summarize_btn = gr.Button("🔍 Tóm tắt văn bản", variant="primary", size="lg")
    
    # Output section
    gr.Markdown("## 📋 Kết quả")
    
    with gr.Row():
        with gr.Column(scale=2):
            summary_output = gr.Textbox(
                label="✅ Bản tóm tắt",
                lines=5,
                max_lines=10,
                show_copy_button=True
            )
        
        with gr.Column(scale=1):
            stats_output = gr.Markdown(label="📊 Thống kê")
    
    status_output = gr.Textbox(label="Trạng thái", visible=False)
    
    # Connect the summarize function
    summarize_btn.click(
        fn=summarize_text,
        inputs=[text_input, max_length, min_length, num_beams],
        outputs=[summary_output, stats_output, status_output]
    )
    
    # Features section
    gr.HTML("""
    <div style="margin-top: 3rem;">
        <h2 style="text-align: center; margin-bottom: 2rem;">🌟 Tính năng nổi bật</h2>
        <div style="display: flex; gap: 1rem; flex-wrap: wrap;">
            <div class="feature-box" style="flex: 1; min-width: 250px;">
                <h4>🎯 Độ chính xác cao</h4>
                <p style="margin: 0;">Mô hình VietAI T5 được huấn luyện chuyên biệt trên dữ liệu tiếng Việt</p>
            </div>
            <div class="feature-box" style="flex: 1; min-width: 250px;">
                <h4>⚡ Xử lý nhanh</h4>
                <p style="margin: 0;">Tóm tắt văn bản dài trong vài giây với độ chính xác cao</p>
            </div>
            <div class="feature-box" style="flex: 1; min-width: 250px;">
                <h4>🔧 Tùy chỉnh linh hoạt</h4>
                <p style="margin: 0;">Điều chỉnh độ dài và chất lượng tóm tắt theo nhu cầu</p>
            </div>
        </div>
    </div>
    """)
    
    # Footer
    gr.HTML("""
    <footer>
        <p>📝 Vietnamese Text Summarization | Powered by VietAI T5 Model</p>
        <p>Built with ❤️ using Gradio and Hugging Face Transformers</p>
    </footer>
    """)

# Launch the app
if __name__ == "__main__":
    app.launch()