import gradio as gr
import os
import ctranslate2
from transformers import M2M100Tokenizer
from huggingface_hub import snapshot_download
import time

# Model settings
MODEL_PATH = os.getenv("MODEL_PATH", "skypro1111/m2m100-ukr-verbalization-ct2")
TOKENIZER_PATH = os.getenv("TOKENIZER_PATH", "skypro1111/m2m100-ukr-verbalization")

# Initialize translator and tokenizer globally
def init_model(model_path=MODEL_PATH, tokenizer_path=TOKENIZER_PATH):
    print("\nInitializing CTranslate2 model and tokenizer...")
    
    # Download the model from HuggingFace Hub
    local_model_path = snapshot_download(
        repo_id=model_path,
        allow_patterns=["*.bin", "*.json", "tokenizer.json", "vocab.json"],
    )
    
    # Initialize translator with optimizations for CPU
    translator = ctranslate2.Translator(
        local_model_path,
        device="cpu",
        compute_type="int8",  # Use INT8 quantization for CPU
        intra_threads=4,      # Adjust based on CPU cores available
    )
    
    # Load tokenizer
    tokenizer = M2M100Tokenizer.from_pretrained(tokenizer_path)
    tokenizer.src_lang = "uk"
    
    return translator, tokenizer

translator, tokenizer = init_model()

def process_text(text: str):
    """Process a single text input using the CTranslate2 model."""
    start_time = time.time()
    
    # Tokenize input
    source = tokenizer.convert_ids_to_tokens(tokenizer.encode(text))
    target_prefix = [tokenizer.lang_code_to_token["uk"]]
    
    # Run inference
    results = translator.translate_batch(
        [source],
        target_prefix=[target_prefix],
        beam_size=1,
        num_hypotheses=1,
        use_vmap=True,
    )
    
    # Get target tokens and decode
    target = results[0].hypotheses[0][1:]  # Remove language token
    output = tokenizer.decode(tokenizer.convert_tokens_to_ids(target))
    
    inference_time = time.time() - start_time
    return output, inference_time

def inference(text):
    """Gradio interface function."""
    output, inference_time = process_text(text)
    return output, f"{inference_time:.3f} секунд"

# Define examples from inference_ct2.py
examples = [
    ["Моя бабуся народилася 07.11.1919, у важкий післявоєнний час."],
    ["Зустріч призначена на 15:30 12.05.2025 у конференц-залі №3."],
    ["Телефонуйте нам за номером +380 (44) 123-45-67 або 0800 500 123."],
    ["Температура повітря сьогодні становить +25°C, а тиск 750 мм.рт.ст."],
    ["ТОВ «Мрія» було засновано 28/06/2022 з початковим капіталом 50 тис. грн."],
    ["Швидкість вітру 15 м/с, видимість 10 км, вологість 65%."],
    ["Потяг №743 Київ-Львів відправляється о 08:45 з платформи №2."],
    ["Ціна на пальне зросла на 2,5 грн/л і становить 54,99 грн."],
    ["Площа квартири 75,5 м², висота стелі 2,75 м."],
    ["Відстань між містами становить 450 км або 280 миль."],
    ["Станом на 7:00 15 квітня 2025 року температура повітря становить +18°C, вологість 60%."]
]

# Define Gradio interface
with gr.Blocks() as interface:
    gr.Markdown("# Вербалізація українського тексту (M2M100-CT2)")
    gr.Markdown("""Модель для перетворення чисел, дат, одиниць вимірювання та інших символьних позначень 
    у їх текстовий запис українською мовою. Використовує оптимізовану CTranslate2 версію для швидкого інференсу.""")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="Введіть текст для вербалізації",
                placeholder="Наприклад: Температура повітря сьогодні становить +25°C"
            )
            submit_btn = gr.Button("Вербалізувати")
            
            gr.Examples(
                examples=examples,
                inputs=input_text,
            )
            
        with gr.Column():
            output_text = gr.Textbox(label="Вербалізований текст")
            inference_time = gr.Textbox(label="Час обробки")
    
    gr.Markdown("""
    ### Можливості моделі:
    - Дати та час
    - Телефонні номери
    - Одиниці вимірювання
    - Грошові суми
    - Числові вирази
    
    ### Технічні деталі:
    - Базова модель: facebook/m2m100_418M
    - Оптимізація: CTranslate2 з INT8 квантизацією
    """)
    
    submit_btn.click(
        fn=inference,
        inputs=input_text,
        outputs=[output_text, inference_time],
    )

# Launch the interface
interface.launch(share=False)