app.py

#!/usr/bin/env python3
"""
Vietnamese Receipt Classification App for Hugging Face Spaces
Complete version with training logging support
"""

import os
import sys
import gradio as gr
import numpy as np
import json
import tempfile
from datetime import datetime
from pathlib import Path
import threading
import time
import io
from PIL import Image
import logging
import markdown
import re

# Add paths for imports
current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, current_dir)
sys.path.insert(0, os.path.join(current_dir, 'src'))

# Google AI Studio imports
try:
    import google.generativeai as genai
    GOOGLE_AI_AVAILABLE = True
except ImportError:
    GOOGLE_AI_AVAILABLE = False
    print("⚠️ Google AI not available. Install: pip install google-generativeai")

# Project imports
try:
    from config import Config
    from src.trainer import ReceiptClassificationTrainer
    from src.utils import predict_samples, preprocess_text_for_prediction
    from src.logger_config import LoggerConfig
    COMPONENTS_AVAILABLE = True
except ImportError as e:
    print(f"⚠️ Project components not available: {e}")
    COMPONENTS_AVAILABLE = False

# ====================================
# LOGGING SETUP FOR TRAINING ONLY
# ====================================
class TrainingLogCapture(logging.Handler):
    """Handler to capture training logs for Gradio display"""
    def __init__(self):
        super().__init__()
        self.logs = []
        self.max_logs = 200  # Increased from 100
        
    def emit(self, record):
        try:
            msg = self.format(record)
            timestamp = datetime.now().strftime('%H:%M:%S')
            log_entry = f"[{timestamp}] {msg}"
            self.logs.append(log_entry)
            # Keep only last max_logs entries to prevent memory issues
            if len(self.logs) > self.max_logs:
                self.logs.pop(0)
        except Exception:
            self.handleError(record)
    
    def get_logs(self, last_n=None):
        """Get last n log entries or all if n is None"""
        if last_n is None:
            return "\n".join(self.logs)
        return "\n".join(self.logs[-last_n:])
    
    def clear_logs(self):
        """Clear all logs"""
        self.logs = []

# Create training log capture instance
training_log_capture = TrainingLogCapture()
training_log_capture.setFormatter(logging.Formatter('%(message)s'))

# ====================================
# GLOBAL VARIABLES
# ====================================
trained_model = None
feature_type = None
vectorizers = None
label_encoder = None
training_status = "Not started"
is_training = False

# ====================================
# GOOGLE AI VISION SETUP
# ====================================
def setup_google_ai():
    """Setup Google AI with API key from environment"""
    if not GOOGLE_AI_AVAILABLE:
        return None
    
    api_key = os.getenv('GOOGLE_AI_API_KEY') or os.getenv('GOOGLE_API_KEY')
    
    if not api_key:
        print("❌ Google AI API key not found in environment variables")
        return None
    
    try:
        genai.configure(api_key=api_key)
        model = genai.GenerativeModel('gemini-1.5-flash')
        print("✅ Google AI Vision model initialized")
        return model
    except Exception as e:
        print(f"❌ Error setting up Google AI: {e}")
        return None

google_vision_model = setup_google_ai()

# ====================================
# TRAINING FUNCTIONS WITH LOGGING
# ====================================
def train_model_background():
    """Train model in background thread with logging"""
    global trained_model, feature_type, vectorizers, label_encoder, training_status, is_training
    
    if not COMPONENTS_AVAILABLE:
        training_status = "❌ Training components not available"
        training_log_capture.logs.append("[ERROR] Training components not available")
        return
    
    try:
        is_training = True
        training_status = "Starting training..."
        
        # Clear previous logs
        training_log_capture.clear_logs()
        
        # Setup training logger with our capture handler
        training_logger = LoggerConfig.setup_training_logger()
        training_logger.addHandler(training_log_capture)
        
        training_logger.info("🚀 Starting training process...")
        print("🚀 Starting training process...")  # Also print
        
        # Check if dataset exists
        if not os.path.exists(Config.DATA_FILE):
            training_status = "Error: Dataset not found"
            training_logger.error(f"Dataset {Config.DATA_FILE} not found")
            print(f"❌ Dataset {Config.DATA_FILE} not found")
            is_training = False
            return
        
        training_status = "Training in progress... (This may take 10-15 minutes)"
        training_logger.info("Training started - this may take 10-15 minutes")
        print("Training started - this may take 10-15 minutes")
        
        # Initialize trainer (will use logging internally)
        trainer = ReceiptClassificationTrainer(Config)
        
        # Add the handler to trainer's logger as well
        if hasattr(trainer, 'logger'):
            trainer.logger.addHandler(training_log_capture)
        
        # Run training pipeline
        best_model, best_feature_type, results = trainer.run_full_pipeline()
        
        # Set global variables
        trained_model = best_model
        feature_type = best_feature_type
        vectorizers = trainer.feature_extractor.get_vectorizers()
        label_encoder = trainer.data_loader.label_encoder
        
        accuracy = results.get('accuracy', 0)
        training_status = f"✅ Training completed! Accuracy: {accuracy:.4f}"
        training_logger.info(f"✅ Training completed with {accuracy:.4f} accuracy")
        print(f"✅ Training completed with {accuracy:.4f} accuracy")
        
    except Exception as e:
        training_status = f"❌ Training failed: {str(e)}"
        training_log_capture.logs.append(f"[ERROR] Training failed: {str(e)}")
        print(f"❌ Training failed: {str(e)}")
    
    finally:
        is_training = False

def get_training_status():
    """Get current training status and logs"""
    # Get all logs for better visibility
    log_text = training_log_capture.get_logs()
    if not log_text:
        log_text = "No logs yet... Click 'Start Training' to begin"
    return training_status, log_text

def start_training():
    """Start training process with logging"""
    global is_training
    
    if not COMPONENTS_AVAILABLE:
        return "❌ Training components not available", "Missing required modules"
    
    if is_training:
        return "⚠️ Training already in progress...", training_log_capture.get_logs()
    
    thread = threading.Thread(target=train_model_background)
    thread.daemon = True
    thread.start()
    
    return "🚀 Training started in background...", "Training initiated... Logs will appear here"

# ====================================
# VISION MODEL FUNCTIONS (NO LOGGING)
# ====================================
def extract_bill_description(image):
    """Extract bill description using Google Vision AI"""
    if not GOOGLE_AI_AVAILABLE or google_vision_model is None:
        return "❌ Google AI Vision không khả dụng. Vui lòng thiết lập GOOGLE_AI_API_KEY hoặc nhập mô tả thủ công."
    
    try:
        if image is None:
            return "❌ Vui lòng upload ảnh hóa đơn"
        
        # Convert image to PIL if needed
        if not isinstance(image, Image.Image):
            image = Image.fromarray(image)
        
        # Prompt for Vietnamese receipt description
        prompt = """
        Bạn là một AI chuyên phân tích hóa đơn Việt Nam. Hãy mô tả chi tiết hóa đơn này theo định dạng sau:

        Mô tả hóa đơn: [Tên cửa hàng/nhà hàng] - [Loại hình kinh doanh] - [Các món/sản phẩm chính] - [Tổng tiền] - [Ngày tháng nếu có] - [Địa điểm nếu có]

        Ví dụ: "Hóa đơn thanh toán tại cửa hàng cà phê Feel Coffee với món Yogurt Very Berry giá 22.000 VND, thanh toán bằng tiền mặt"

        Hãy mô tả hóa đơn trong ảnh theo format tương tự, bằng tiếng Việt:
        """
        
        # Generate description
        response = google_vision_model.generate_content([prompt, image])
        description = response.text.strip()
        
        if description:
            return description
        else:
            return "❌ Không thể trích xuất thông tin từ ảnh. Vui lòng thử ảnh khác hoặc nhập mô tả thủ công."
            
    except Exception as e:
        return f"❌ Lỗi khi phân tích ảnh: {str(e)}"

def process_image_and_extract(image):
    """Process uploaded image and extract description"""
    if image is None:
        return "Vui lòng upload ảnh hóa đơn"
    
    description = extract_bill_description(image)
    return description

# function load and convert README
def load_readme():
    """Load and convert README.md to HTML for display"""
    try:
        with open("README.md", "r", encoding="utf-8") as file:
            readme_content = file.read()
            
        # Remove HF metadata header (between ---)
        readme_content = re.sub(r'^---\n.*?\n---\n', '', readme_content, flags=re.DOTALL)
        
        # Convert markdown to HTML
        html_content = markdown.markdown(
            readme_content,
            extensions=[
                'markdown.extensions.tables',
                'markdown.extensions.fenced_code',
                'markdown.extensions.codehilite',
                'markdown.extensions.toc',
                'markdown.extensions.nl2br'
            ]
        )
        
        # Add custom CSS for better styling
        styled_html = f"""
        <div style="padding: 20px; max-width: 1200px; margin: 0 auto;">
            <style>
                /* General styles */
                h1 {{ color: #2c3e50; border-bottom: 3px solid #3498db; padding-bottom: 10px; }}
                h2 {{ color: #34495e; margin-top: 30px; border-bottom: 2px solid #ecf0f1; padding-bottom: 8px; }}
                h3 {{ color: #7f8c8d; margin-top: 20px; }}
                
                /* Table styles */
                table {{
                    border-collapse: collapse;
                    width: 100%;
                    margin: 20px 0;
                    box-shadow: 0 2px 3px rgba(0,0,0,0.1);
                }}
                th {{
                    background: linear-gradient(90deg, #667eea 0%, #764ba2 100%);
                    color: white;
                    padding: 12px;
                    text-align: left;
                    font-weight: bold;
                }}
                td {{
                    padding: 10px;
                    border-bottom: 1px solid #ecf0f1;
                }}
                tr:hover {{
                    background-color: #f8f9fa;
                }}
                
                /* Code block styles */
                pre {{
                    background-color: #f8f9fa;
                    color: #212529;
                    padding: 15px;
                    border-radius: 8px;
                    overflow-x: auto;
                    margin: 15px 0;
                }}
                code {{
                    background-color: #ecf0f1;
                    padding: 2px 6px;
                    border-radius: 3px;
                    font-family: 'Courier New', monospace;
                }}
                pre code {{
                    background-color: transparent;
                    padding: 0;
                }}
                
                /* List styles */
                ul, ol {{
                    margin: 15px 0;
                    padding-left: 30px;
                }}
                li {{
                    margin: 8px 0;
                    line-height: 1.6;
                }}
                
                /* Link styles */
                a {{
                    color: #3498db;
                    text-decoration: none;
                    transition: color 0.3s;
                }}
                a:hover {{
                    color: #2980b9;
                    text-decoration: underline;
                }}
                
                /* Blockquote styles */
                blockquote {{
                    border-left: 4px solid #3498db;
                    padding-left: 20px;
                    margin: 20px 0;
                    color: #7f8c8d;
                    font-style: italic;
                }}
                
                /* Horizontal rule */
                hr {{
                    border: none;
                    height: 2px;
                    background: linear-gradient(90deg, transparent, #bdc3c7, transparent);
                    margin: 30px 0;
                }}
                
                /* Badge styles */
                img[alt*="badge"] {{
                    margin: 0 5px;
                }}
                
                /* Emoji support */
                .emoji {{
                    font-size: 1.2em;
                    margin: 0 3px;
                }}
            </style>
            {html_content}
        </div>
        """
        
        return styled_html
        
    except FileNotFoundError:
        return """
        <div style="padding: 20px; text-align: center;">
            <h2 style="color: #e74c3c;">❌ README.md not found</h2>
            <p>Please ensure README.md file exists in the root directory.</p>
        </div>
        """
    except Exception as e:
        return f"""
        <div style="padding: 20px; text-align: center;">
            <h2 style="color: #e74c3c;">❌ Error loading README</h2>
            <p>Error: {str(e)}</p>
        </div>
        """
# ====================================
# PREDICTION FUNCTIONS (NO LOGGING)
# ====================================
def predict_bill_class(description):
    """Predict bill class from description"""
    global trained_model, feature_type, vectorizers, label_encoder
    
    if not COMPONENTS_AVAILABLE:
        return "❌ Prediction components not available", "", "Components missing"
    
    if trained_model is None:
        return "❌ Model chưa được train. Vui lòng đợi quá trình training hoàn tất.", "", "Model not ready"
    
    if not description or description.strip() == "":
        return "❌ Vui lòng nhập mô tả hóa đơn", "", "Empty description"
    
    try:
        # Predict
        predictions, probabilities = predict_samples(
            [description], trained_model, feature_type, vectorizers, label_encoder
        )
        
        predicted_class = predictions[0]
        confidence = max(probabilities[0])
        
        # Get top 3 predictions
        top_3_indices = np.argsort(probabilities[0])[-3:][::-1]
        top_3_results = []
        
        for i, idx in enumerate(top_3_indices, 1):
            label = label_encoder.classes_[idx]
            conf = probabilities[0][idx]
            top_3_results.append(f"{i}. {label}: {conf:.3f}")
        
        result_text = f"🎯 Dự đoán: {predicted_class}\n📊 Độ tin cậy: {confidence:.3f}"
        top_3_text = "📊 Top 3 dự đoán:\n" + "\n".join(top_3_results)
        status = f"✅ Đã phân loại thành công với độ tin cậy {confidence:.1%}"
        
        return result_text, top_3_text, status
        
    except Exception as e:
        return f"❌ Lỗi khi dự đoán: {str(e)}", "", f"Error: {str(e)}"

def predict_from_image_and_text(image, manual_description):
    """Combined prediction from image and manual text"""
    
    # Use manual description if provided, otherwise extract from image
    if manual_description and manual_description.strip():
        description = manual_description.strip()
        source_info = "📝 Sử dụng mô tả thủ công"
    elif image is not None:
        description = extract_bill_description(image)
        source_info = "🖼️ Trích xuất từ ảnh"
        
        # Check if extraction failed
        if description.startswith("❌"):
            return description, "", description, description
    else:
        return "❌ Vui lòng upload ảnh hoặc nhập mô tả thủ công", "", "No input provided", ""
    
    # Make prediction
    result, top_3, status = predict_bill_class(description)
    
    # Prepare full description info
    full_description = f"{source_info}\n\n📄 Mô tả hóa đơn:\n{description}"
    
    return result, top_3, status, full_description

# ====================================
# GRADIO INTERFACE
# ====================================
def create_interface():
    """Create Gradio interface with training logging only"""
    
    # Custom CSS for scrollable log
    css = """
    .gradio-container {
        max-width: 1200px !important;
    }
    .main-header {
        text-align: center;
        margin: 20px 0;
        padding: 20px;
        background: linear-gradient(90deg, #667eea 0%, #764ba2 100%);
        color: white;
        border-radius: 10px;
    }
    /* Make the log textarea scrollable */
    textarea {
        overflow-y: auto !important;
        font-family: 'Courier New', monospace;
        font-size: 12px;
    }
    """
    
    with gr.Blocks(title="Vietnamese Receipt Classification", css=css) as interface:
        
        # Header
        gr.HTML("""
        <div class="main-header">
            <h1>🧾 Vietnamese Receipt Classification</h1>
            <p>Ứng dụng phân loại hóa đơn Việt Nam sử dụng GA-optimized Ensemble + Google AI Vision</p>
        </div>
        """)
        
        with gr.Tabs():
            
            # ====================================
            # TAB 1: MODEL TRAINING
            # ====================================
            with gr.Tab("🚀 Model Training"):
                
                gr.HTML("<h3>🏋️ Training Management</h3>")
                
                with gr.Row():
                    train_btn = gr.Button("🚀 Start Training", variant="primary", size="lg")
                    refresh_btn = gr.Button("🔄 Refresh Status", variant="secondary")
                
                status_display = gr.Textbox(
                    label="📊 Training Status",
                    value="Click 'Start Training' to begin",
                    interactive=False,
                    lines=2
                )
                
                # Increased lines and set max_lines for scrollability
                log_display = gr.Textbox(
                    label="📝 Training Log (Scrollable)",
                    lines=20,  # Increased from 10
                    max_lines=20,  # Set max lines for scrolling
                    interactive=False,
                    placeholder="Training logs will appear here...",
                    autoscroll=True  # Auto scroll to bottom
                )
                
                # Training info
                gr.HTML("""
                <div style="margin-top: 20px; padding: 20px; background-color: #f8f9fa; border-radius: 8px; border-left: 4px solid #007bff;">
                    <h4>📋 Training Information</h4>
                    <ul style="margin: 10px 0; padding-left: 20px;">
                        <li><strong>Algorithm:</strong> GA-optimized Voting Ensemble (KNN + Decision Tree + Naive Bayes)</li>
                        <li><strong>Features:</strong> BoW, TF-IDF, Sentence Embeddings (all-MiniLM-L6-v2)</li>
                        <li><strong>Optimization:</strong> Genetic Algorithm (Population: 30, Generations: 15)</li>
                        <li><strong>Evaluation:</strong> 3-fold Cross-Validation</li>
                        <li><strong>Expected Time:</strong> 10-15 minutes on free tier</li>
                        <li><strong>Expected Accuracy:</strong> 85-95% depending on dataset quality</li>
                        <li><strong>Logging:</strong> All outputs are captured in scrollable log above</li>
                        <li><strong>Refresh:</strong> Click refresh button to update logs during training</li>
                    </ul>
                </div>
                """)
                
                # Event handlers for training tab
                train_btn.click(fn=start_training, outputs=[status_display, log_display])
                refresh_btn.click(fn=get_training_status, outputs=[status_display, log_display])
            
            # ====================================
            # TAB 2: BILL CLASSIFICATION
            # ====================================
            with gr.Tab("🔮 Bill Classification"):
                
                gr.HTML("<h3>🎯 Phân loại hóa đơn từ ảnh hoặc text</h3>")
                
                with gr.Row():
                    # Left column - Input
                    with gr.Column(scale=1):
                        gr.HTML("<h4>📸 Upload ảnh hóa đơn</h4>")
                        
                        image_input = gr.Image(
                            label="Ảnh hóa đơn",
                            type="pil",
                            height=250
                        )
                        
                        extract_btn = gr.Button("🔍 Trích xuất mô tả từ ảnh", variant="secondary")
                        
                        gr.HTML("<h4>📝 Hoặc nhập mô tả thủ công</h4>")
                        
                        manual_input = gr.Textbox(
                            label="Mô tả hóa đơn",
                            placeholder="Ví dụ: Hóa đơn thanh toán tại cửa hàng cà phê Feel Coffee với món Yogurt Very Berry giá 22.000 VND",
                            lines=4
                        )
                        
                        predict_btn = gr.Button("🎯 Dự đoán phân loại", variant="primary", size="lg")
                    
                    # Right column - Output
                    with gr.Column(scale=1):
                        gr.HTML("<h4>📄 Thông tin đã xử lý</h4>")
                        
                        processed_info = gr.Textbox(
                            label="Nguồn và mô tả",
                            lines=6,
                            interactive=False
                        )
                        
                        gr.HTML("<h4>🎯 Kết quả phân loại</h4>")
                        
                        result_display = gr.Textbox(
                            label="Dự đoán chính",
                            lines=3,
                            interactive=False
                        )
                        
                        top3_display = gr.Textbox(
                            label="Top 3 dự đoán",
                            lines=4,
                            interactive=False
                        )
                        
                        status_output = gr.Textbox(
                            label="Trạng thái",
                            lines=2,
                            interactive=False
                        )
                
                # Examples section
                gr.HTML("""
                <div style="margin-top: 20px; padding: 15px; background-color: #e8f4fd; border-radius: 8px;">
                    <h4>💡 Ví dụ các loại hóa đơn</h4>
                    <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 20px; margin-top: 10px;">
                        <div>
                            <ul style="margin: 0; padding-left: 20px;">
                                <li><strong>Ăn uống ngoài hàng:</strong> Nhà hàng, quán cà phê, fast food</li>
                                <li><strong>Siêu thị tổng hợp:</strong> VinMart, Co.opMart, Big C, Lotte</li>
                            </ul>
                        </div>
                        <div>
                            <ul style="margin: 0; padding-left: 20px;">
                                <li><strong>Sữa & Đồ uống:</strong> Sữa, nước ngọt, đồ uống các loại</li>
                                <li><strong>Tiện ích:</strong> Điện, nước, internet, di động</li>
                            </ul>
                        </div>
                    </div>
                </div>
                """)
                
                # Event handlers for classification tab
                extract_btn.click(
                    fn=process_image_and_extract,
                    inputs=[image_input],
                    outputs=[manual_input]
                )
                
                predict_btn.click(
                    fn=predict_from_image_and_text,
                    inputs=[image_input, manual_input],
                    outputs=[result_display, top3_display, status_output, processed_info]
                )
            
            # ====================================
            # TAB 3: ABOUT & HELP
            # ====================================
            with gr.Tab("ℹ️ About & Help"):
                
                gr.HTML("""
                <div style="padding: 20px;">
                    <h2 style="color: #2c3e50;">🧾 Vietnamese Receipt Classification System</h2>
                    
                    <div class="info-section">
                        <h3>🎯 Tính năng chính</h3>
                        <ul>
                            <li><strong>🤖 AI Vision:</strong> Trích xuất mô tả từ ảnh hóa đơn bằng Google Gemini Vision API</li>
                            <li><strong>🧬 GA Optimization:</strong> Tối ưu hóa ensemble classifier bằng Genetic Algorithm</li>
                            <li><strong>📊 Multi-feature:</strong> Kết hợp BoW, TF-IDF và Sentence Embeddings</li>
                            <li><strong>🗳️ Voting Ensemble:</strong> KNN + Decision Tree + Naive Bayes với trọng số tối ưu</li>
                            <li><strong>⚡ Real-time:</strong> Training và prediction trực tiếp trên web</li>
                        </ul>
                    </div>
                    
                    <div class="example-section">
                        <h3>🔧 Công nghệ sử dụng</h3>
                        <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 15px;">
                            <div>
                                <h4 style="color: #0d47a1;">Machine Learning:</h4>
                                <ul style="color: #1565c0;">
                                    <li>scikit-learn</li>
                                    <li>sentence-transformers</li>
                                    <li>DEAP (Genetic Algorithm)</li>
                                </ul>
                            </div>
                            <div>
                                <h4 style="color: #0d47a1;">AI Vision:</h4>
                                <ul style="color: #1565c0;">
                                    <li>Google Gemini Vision</li>
                                    <li>PIL (Image Processing)</li>
                                    <li>Gradio Interface</li>
                                </ul>
                            </div>
                        </div>
                    </div>
                    
                    <div class="success-section">
                        <h3>🚀 Hướng dẫn sử dụng</h3>
                        <ol style="color: #155724;">
                            <li><strong>Training:</strong> Bắt đầu với tab "🚀 Model Training", click "Start Training" và đợi 10-15 phút</li>
                            <li><strong>Monitor:</strong> Click "Refresh Status" để cập nhật logs trong quá trình training</li>
                            <li><strong>Classification:</strong> Chuyển sang tab "🔮 Bill Classification"</li>
                            <li><strong>Upload ảnh:</strong> Kéo thả ảnh hóa đơn vào khung "Upload ảnh hóa đơn"</li>
                            <li><strong>Extract text:</strong> Click "🔍 Trích xuất mô tả từ ảnh" (cần Google AI API key)</li>
                            <li><strong>Manual input:</strong> Hoặc nhập mô tả thủ công vào text box</li>
                            <li><strong>Predict:</strong> Click "🎯 Dự đoán phân loại" để xem kết quả</li>
                            <li><strong>Results:</strong> Xem dự đoán chính + top 3 alternatives với confidence scores</li>
                        </ol>
                    </div>
                    
                    <div class="warning-section">
                        <h3>⚠️ Lưu ý quan trọng</h3>
                        <ul style="color: #856404;">
                            <li><strong>Google AI API:</strong> Để sử dụng tính năng trích xuất từ ảnh, cần thiết lập GOOGLE_AI_API_KEY trong environment variables</li>
                            <li><strong>Dataset:</strong> App cần file viet_receipt_categorized_label.xlsx để training</li>
                            <li><strong>Memory:</strong> Training có thể tốn nhiều RAM, nên dùng trên máy có đủ bộ nhớ</li>
                            <li><strong>Time:</strong> Quá trình training mất 10-15 phút, vui lòng kiên nhẫn</li>
                            <li><strong>Logs:</strong> Training log có thể scroll để xem toàn bộ quá trình</li>
                        </ul>
                    </div>
                    
                    <div style="text-align: center; margin-top: 30px; padding: 20px; background: linear-gradient(45deg, #2c3e50, #3498db); color: white; border-radius: 8px;">
                        <h3>🎉 Developed with ❤️ for Vietnamese NLP Community</h3>
                        <p>Powered by Hugging Face 🤗 | Google AI Studio | Gradio</p>
                    </div>
                </div>
                """)
            with gr.Tab("📚 Documentation"):
                gr.HTML("<h3>📖 Complete Project Documentation</h3>")
                
                # Refresh button để reload README
                with gr.Row():
                    refresh_docs_btn = gr.Button(
                        "🔄 Refresh Documentation", 
                        variant="secondary",
                        size="sm"
                    )
                    
                    # Search box cho documentation
                    search_box = gr.Textbox(
                        placeholder="🔍 Search in documentation...",
                        label="Search",
                        scale=3
                    )
                
                # README content display
                readme_display = gr.HTML(
                    value=load_readme(),
                    label="README Documentation"
                )
                
                # JavaScript for search functionality
                gr.HTML("""
                <script>
                function searchInDocs() {
                    const searchTerm = document.querySelector('input[placeholder*="Search in documentation"]').value.toLowerCase();
                    const content = document.querySelector('[label="README Documentation"]');
                    
                    if (!searchTerm) {
                        // Remove all highlights if search is empty
                        content.innerHTML = content.innerHTML.replace(/<mark[^>]*>(.*?)<\/mark>/gi, '$1');
                        return;
                    }
                    
                    // Remove previous highlights
                    content.innerHTML = content.innerHTML.replace(/<mark[^>]*>(.*?)<\/mark>/gi, '$1');
                    
                    // Add new highlights
                    const regex = new RegExp(`(${searchTerm})`, 'gi');
                    content.innerHTML = content.innerHTML.replace(regex, '<mark style="background-color: yellow; padding: 2px;">$1</mark>');
                    
                    // Scroll to first match
                    const firstMatch = content.querySelector('mark');
                    if (firstMatch) {
                        firstMatch.scrollIntoView({ behavior: 'smooth', block: 'center' });
                    }
                }
                
                // Add event listener when page loads
                document.addEventListener('DOMContentLoaded', function() {
                    const searchInput = document.querySelector('input[placeholder*="Search in documentation"]');
                    if (searchInput) {
                        searchInput.addEventListener('input', searchInDocs);
                    }
                });
                </script>
                """)
                
                # Quick navigation links
                gr.HTML("""
                <div style="margin-top: 20px; padding: 15px; background-color: #f8f9fa; border-radius: 8px;">
                    <h4>⚡ Quick Links</h4>
                    <div style="display: flex; flex-wrap: wrap; gap: 10px; margin-top: 10px;">
                        <a href="#overview" style="padding: 5px 15px; background: #3498db; color: white; border-radius: 5px; text-decoration: none;">Overview</a>
                        <a href="#quick-deployment-guide" style="padding: 5px 15px; background: #2ecc71; color: white; border-radius: 5px; text-decoration: none;">Deployment</a>
                        <a href="#user-guide" style="padding: 5px 15px; background: #e74c3c; color: white; border-radius: 5px; text-decoration: none;">User Guide</a>
                        <a href="#technical-architecture" style="padding: 5px 15px; background: #9b59b6; color: white; border-radius: 5px; text-decoration: none;">Technical</a>
                        <a href="#troubleshooting" style="padding: 5px 15px; background: #f39c12; color: white; border-radius: 5px; text-decoration: none;">Troubleshooting</a>
                        <a href="#requirements" style="padding: 5px 15px; background: #34495e; color: white; border-radius: 5px; text-decoration: none;">Requirements</a>
                    </div>
                </div>
                """)
                
                # Event handler for refresh button
                refresh_docs_btn.click(
                    fn=load_readme,
                    outputs=[readme_display]
                )
                
                # Optional: Add a download button for README
                gr.HTML("""
                <div style="margin-top: 20px; text-align: center;">
                    <a href="README.md" download="README.md" 
                       style="display: inline-block; padding: 10px 20px; background: linear-gradient(90deg, #667eea 0%, #764ba2 100%); 
                              color: white; border-radius: 5px; text-decoration: none; font-weight: bold;">
                        📥 Download README.md
                    </a>
                </div>
                """)
        
        # Load initial status when interface starts
        interface.load(fn=get_training_status, outputs=[status_display, log_display])
    
    return interface

# ====================================
# MAIN APPLICATION
# ====================================
if __name__ == "__main__":
    print("🚀 Starting Vietnamese Receipt Classification App...")
    print("="*60)
    
    # Check dependencies
    print("📋 Checking dependencies...")
    
    if COMPONENTS_AVAILABLE:
        print("✅ Project components: Ready")
        
        # Check dataset
        try:
            if os.path.exists(Config.DATA_FILE):
                print(f"✅ Dataset: Found {Config.DATA_FILE}")
            else:
                print(f"⚠️ Dataset: {Config.DATA_FILE} not found")
        except:
            print("⚠️ Config not available")
    else:
        print("⚠️ Project components: Not available")
    
    if GOOGLE_AI_AVAILABLE and google_vision_model is not None:
        print("✅ Google AI Vision: Ready")
    else:
        print("⚠️ Google AI Vision: Not available")
        print("   💡 Set GOOGLE_AI_API_KEY environment variable to enable")
    
    print("🎨 Creating Gradio interface...")
    app = create_interface()
    
    print("🌐 Launching app...")
    print("="*60)
    
    # Launch the app
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True
    )