# -*- coding: utf-8 -*-
"""
Ultra-FineWeb Classifier - Hugging Face Space Demo
A lightweight fastText-based classifier for filtering high-quality web data.
"""

import os
import re
import unicodedata
from typing import Tuple

import gradio as gr
from huggingface_hub import hf_hub_download

# Lazy loading for heavy dependencies
_tokenizer = None
_fasttext_models = {}

MODEL_REPO = "openbmb/Ultra-FineWeb-classifier"


def get_tokenizer():
    """Lazy load tokenizer."""
    global _tokenizer
    if _tokenizer is None:
        from transformers import AutoTokenizer
        
        # Download tokenizer files from the model repo
        tokenizer_path = hf_hub_download(
            repo_id=MODEL_REPO,
            filename="local_tokenizer/tokenizer.json",
            local_dir="./model_cache",
        )
        tokenizer_dir = os.path.dirname(tokenizer_path)
        
        # Download other tokenizer files
        for filename in [
            "local_tokenizer/tokenizer_config.json",
            "local_tokenizer/special_tokens_map.json",
        ]:
            hf_hub_download(
                repo_id=MODEL_REPO,
                filename=filename,
                local_dir="./model_cache",
            )
        
        _tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir)
    return _tokenizer


def get_fasttext_model(language: str):
    """Lazy load fastText model for specific language."""
    global _fasttext_models
    
    if language not in _fasttext_models:
        import fasttext
        
        model_filename = f"classifiers/ultra_fineweb_{language}.bin"
        model_path = hf_hub_download(
            repo_id=MODEL_REPO,
            filename=model_filename,
            local_dir="./model_cache",
        )
        _fasttext_models[language] = fasttext.load_model(model_path)
    
    return _fasttext_models[language]


def fasttext_preprocess(content: str, tokenizer) -> str:
    """
    Preprocess content for fastText inference.
    
    Steps:
    1. Remove multiple newlines
    2. Lowercase
    3. Remove diacritics
    4. Word segmentation using tokenizer
    5. Handle escape characters
    """
    # 1. Remove multiple newlines
    content = re.sub(r'\n{3,}', '\n\n', content)
    
    # 2. Lowercase
    content = content.lower()
    
    # 3. Remove diacritics
    content = ''.join(
        c for c in unicodedata.normalize('NFKD', content)
        if unicodedata.category(c) != 'Mn'
    )
    
    # 4. Word segmentation
    token_ids = tokenizer.encode(content, add_special_tokens=False)
    single_text_list = []
    for token_id in token_ids:
        curr_text = tokenizer.decode([token_id])
        single_text_list.append(curr_text)
    
    content = ' '.join(single_text_list)
    
    # 5. Handle escape characters
    content = re.sub(r'\n', '\\\\n', content)
    content = re.sub(r'\r', '\\\\r', content)
    content = re.sub(r'\t', '\\\\t', content)
    content = re.sub(r' +', ' ', content)
    content = content.strip()
    
    return content


def fasttext_infer(norm_content: str, fasttext_model) -> Tuple[str, float]:
    """
    Run fastText inference.
    
    Returns:
        Tuple of (label, score) where score is the probability of being high-quality.
    """
    pred_label, pred_prob = fasttext_model.predict(norm_content)
    pred_label = pred_label[0]
    score = min(pred_prob.tolist()[0], 1.0)
    
    # Convert to positive score (probability of being high-quality)
    if pred_label == "__label__neg":
        score = 1 - score
    
    return pred_label, score


def classify_text(content: str, language: str) -> Tuple[str, str]:
    """
    Main classification function.
    
    Args:
        content: Text to classify
        language: Language code ("en" or "zh")
    
    Returns:
        Tuple of (pred_label, score_display)
    """
    if not content or not content.strip():
        return "N/A", "N/A"
    
    try:
        # Get tokenizer and model
        tokenizer = get_tokenizer()
        fasttext_model = get_fasttext_model(language)
        
        # Preprocess
        norm_content = fasttext_preprocess(content, tokenizer)
        
        # Inference
        pred_label, score = fasttext_infer(norm_content, fasttext_model)
        
        score_display = f"{score:.6f}"
        
        return pred_label, score_display
        
    except Exception as e:
        return "Error", str(e)


# Example texts
EXAMPLE_EN = """Machine learning is a subset of artificial intelligence that enables systems to learn and improve from experience without being explicitly programmed. It focuses on developing computer programs that can access data and use it to learn for themselves.

The process begins with observations or data, such as examples, direct experience, or instruction, in order to look for patterns in data and make better decisions in the future based on the examples that we provide."""

EXAMPLE_ZH = """机器学习是人工智能的一个子集，它使系统能够从经验中学习和改进，而无需显式编程。它专注于开发能够访问数据并使用数据自行学习的计算机程序。

这个过程从观察或数据开始，例如示例、直接经验或指令，以便在数据中寻找模式，并根据我们提供的示例在未来做出更好的决策。"""


# Custom CSS
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap');

:root {
    --bg: #f5f7fb;
    --card: #ffffff;
    --text: #0f172a;
    --muted: #6b7280;
    --border: #e5e7eb;
    --primary: #5b5ce2;
    --primary-600: #4f46e5;
    --shadow: 0 10px 30px rgba(15, 23, 42, 0.08);
}

.gradio-container {
    font-family: 'Inter', system-ui, -apple-system, sans-serif !important;
    background: var(--bg) !important;
    min-height: 100vh;
    padding: 16px !important;
    --button-primary-background-fill: var(--primary);
    --button-primary-background-fill-hover: var(--primary-600);
    --button-primary-border-color: var(--primary);
    --button-primary-border-color-hover: var(--primary-600);
    --button-primary-text-color: #ffffff;
    --button-primary-text-color-hover: #ffffff;
    --button-primary-shadow: none;
    --button-primary-shadow-hover: none;
    --button-primary-shadow-active: none;
    --button-secondary-background-fill: #ffffff;
    --button-secondary-background-fill-hover: #f8fafc;
    --button-secondary-border-color: #cbd5e1;
    --button-secondary-border-color-hover: #94a3b8;
    --button-secondary-text-color: #475569;
    --button-secondary-text-color-hover: #0f172a;
    --button-secondary-shadow: none;
    --button-secondary-shadow-hover: none;
    --button-secondary-shadow-active: none;
    --checkbox-border-width: 1px;
    --checkbox-border-color: #cbd5e1;
    --checkbox-border-color-hover: #a5b4fc;
    --checkbox-border-color-focus: #818cf8;
    --checkbox-border-color-selected: var(--primary);
    --checkbox-background-color: #ffffff;
    --checkbox-background-color-hover: #eef2ff;
    --checkbox-background-color-focus: #e0e7ff;
    --checkbox-background-color-selected: var(--primary);
    --checkbox-shadow: none;
}

.main-title {
    color: var(--primary) !important;
    font-weight: 700 !important;
    font-size: 2.2rem !important;
    text-align: center !important;
    margin-bottom: 0.25rem !important;
    letter-spacing: -0.01em !important;
}

.subtitle {
    text-align: center !important;
    color: var(--muted) !important;
    font-size: 1rem !important;
    margin-bottom: 2rem !important;
    font-weight: 400 !important;
}

.gr-box {
    border-radius: 16px !important;
    border: 1px solid var(--border) !important;
    background: var(--card) !important;
    box-shadow: var(--shadow) !important;
}

.section-header {
    color: var(--text) !important;
    font-weight: 600 !important;
    font-size: 1rem !important;
    line-height: 1.1 !important;
    margin-bottom: 0.4rem !important;
}

.gr-input, .gr-textarea, .gr-textbox {
    background: #f9fafb !important;
    border: 1px solid var(--border) !important;
    border-radius: 10px !important;
    color: var(--text) !important;
    font-size: 0.95rem !important;
}

.gr-input:focus, .gr-textarea:focus, .gr-textbox:focus {
    border-color: #c7d2fe !important;
    box-shadow: 0 0 0 3px rgba(99, 102, 241, 0.15) !important;
}

.gr-button-primary {
    background: var(--primary) !important;
    border: none !important;
    font-weight: 600 !important;
    font-size: 1rem !important;
    padding: 12px 20px !important;
    border-radius: 10px !important;
    color: #ffffff !important;
    transition: background 0.2s ease !important;
}

.gr-button-primary:hover {
    background: var(--primary-600) !important;
}

button.primary {
    background: var(--primary) !important;
    border-color: var(--primary) !important;
}

button.primary:hover {
    background: var(--primary-600) !important;
    border-color: var(--primary-600) !important;
}

.gr-button-secondary {
    background: #ffffff !important;
    border: 1px solid #cbd5e1 !important;
    color: #475569 !important;
    font-weight: 500 !important;
    border-radius: 10px !important;
}

.example-buttons {
    display: flex !important;
    gap: 12px !important;
}

.example-buttons > * {
    flex: 1 1 0 !important;
}

.example-btn button {
    width: 100% !important;
    display: flex !important;
    align-items: center !important;
    justify-content: center !important;
    background: #ffffff !important;
    border: 2px solid #cbd5e1 !important;
    color: #334155 !important;
    font-weight: 600 !important;
    border-radius: 10px !important;
    padding: 10px 14px !important;
    box-shadow: 0 1px 2px rgba(15, 23, 42, 0.06) !important;
}

.example-btn button:hover {
    background: #f8fafc !important;
    border-color: #94a3b8 !important;
}

label {
    color: var(--muted) !important;
    font-weight: 500 !important;
}

input[type="radio"] {
    accent-color: var(--primary) !important;
}

.gr-markdown {
    color: var(--text) !important;
}

.gr-markdown strong {
    color: var(--primary-600) !important;
}

.app-footer {
    text-align: center;
    margin-top: 2rem;
    padding: 1.25rem;
    color: var(--muted);
    font-size: 0.9rem;
    border-top: 1px solid var(--border);
}

.app-footer a {
    color: var(--primary-600);
    text-decoration: none;
}

/* Loading logo tint (Gradio/HF) */
gradio-app img[src*="logo"],
gradio-app img[src*="gradio"],
gradio-app img[alt*="logo" i],
gradio-app svg[aria-label*="logo" i],
gradio-app svg[role="img"] {
    filter: hue-rotate(235deg) saturate(1.4) brightness(0.95);
}

footer {
    display: none !important;
}
"""

# Build Gradio interface
with gr.Blocks(title="UltraFineWeb-L2-Selector", css=custom_css) as demo:
    gr.HTML('<h1 class="main-title">UltraFineWeb-L2-Selector</h1>')
    gr.HTML('<p class="subtitle">Lightweight fastText-based classifier for high-quality web data filtering</p>')
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.HTML('<div class="section-header">Input</div>')
            
            language = gr.Radio(
                choices=[("English", "en"), ("中文", "zh")],
                value="en",
                label="Language / 语言",
                info="Select the language of your content",
            )
            
            content_input = gr.Textbox(
                label="Content to Classify",
                placeholder="Paste your text content here...",
                lines=12,
                max_lines=20,
                value=EXAMPLE_EN,
            )
            
            with gr.Row():
                classify_btn = gr.Button("Classify", variant="primary", size="lg")
                clear_btn = gr.Button("Clear", variant="secondary", size="lg")
            
            # Example texts section removed per request.
        
        with gr.Column(scale=1):
            gr.HTML('<div class="section-header">Output</div>')
            
            label_output = gr.Textbox(
                label="Predicted Label",
                interactive=False,
            )
            score_output = gr.Textbox(
                label="Score",
                interactive=False,
            )
    
    # Event handlers
    classify_btn.click(
        fn=classify_text,
        inputs=[content_input, language],
        outputs=[label_output, score_output],
    )
    
    def clear_all():
        return "", "en", "", ""
    
    clear_btn.click(
        fn=clear_all,
        outputs=[content_input, language, label_output, score_output],
    )
    
    # Auto-update example when language changes
    def update_example_on_language_change(lang):
        if lang == "zh":
            return EXAMPLE_ZH
        return EXAMPLE_EN
    
    language.change(
        fn=update_example_on_language_change,
        inputs=[language],
        outputs=[content_input],
    )
    
    # Footer
    gr.HTML("""
    <div class="app-footer">
        <p><strong>Ultra-FineWeb Classifier</strong> - Part of the <a href="https://huggingface.co/openbmb/Ultra-FineWeb-classifier" target="_blank">Ultra-FineWeb</a> Project</p>
        <p>Based on fastText for efficient web data quality classification. Supports English and Chinese.</p>
        <p><a href="https://arxiv.org/abs/2505.05427" target="_blank">Technical Report</a> | <a href="https://huggingface.co/datasets/openbmb/Ultra-FineWeb" target="_blank">Dataset</a></p>
    </div>
    """)


if __name__ == "__main__":
    demo.launch()