import gradio as gr
from transformers import MarianMTModel, MarianTokenizer
import torch

# Limitar threads de CPU para optimizar
torch.set_num_threads(4)

# Diccionario de modelos multilenguaje OPUS-MT (origen -> destino)
modelos = {
    ("mul", "en"): "Helsinki-NLP/opus-mt-mul-en",
    ("en", "es"): "Helsinki-NLP/opus-mt-en-es",
    ("es", "en"): "Helsinki-NLP/opus-mt-es-en",
    ("en", "fr"): "Helsinki-NLP/opus-mt-en-fr",
    ("fr", "en"): "Helsinki-NLP/opus-mt-fr-en",
}

# Cache de modelos cargados
modelo_cache = {}

# Función para cargar modelo y tokenizer según idioma (caché)
def cargar_modelo(origen, destino):
    key = (origen, destino)
    if key in modelo_cache:
        return modelo_cache[key]
    model_name = modelos.get(key, "Helsinki-NLP/opus-mt-mul-en")
    tokenizer = MarianTokenizer.from_pretrained(model_name)
    model = MarianMTModel.from_pretrained(model_name)
    modelo_cache[key] = (tokenizer, model)
    return tokenizer, model

# Dividir texto en fragmentos ≤512 tokens
def dividir_texto(texto, max_len=512):
    palabras = texto.split()
    fragmentos = []
    fragmento = []
    for palabra in palabras:
        fragmento.append(palabra)
        if len(fragmento) >= max_len:
            fragmentos.append(" ".join(fragmento))
            fragmento = []
    if fragmento:
        fragmentos.append(" ".join(fragmento))
    return fragmentos

# Función de traducción optimizada para textos largos
def traducir_fragmentos(texto, idioma_origen, idioma_destino):
    tokenizer, model = cargar_modelo(idioma_origen, idioma_destino)
    fragmentos = dividir_texto(texto, max_len=512)
    traduccion = []
    for frag in fragmentos:
        inputs = tokenizer(frag, return_tensors="pt", padding=True, truncation=True)
        translated = model.generate(**inputs, max_length=512)
        traduccion.append(tokenizer.decode(translated[0], skip_special_tokens=True))
    return " ".join(traduccion)

# Lista de idiomas disponibles
idiomas = ["mul", "en", "es", "fr"]

# Crear interfaz Gradio
iface = gr.Interface(
    fn=traducir_fragmentos,
    inputs=[
        gr.Textbox(label="Texto a traducir", lines=8, placeholder="Escribe aquí..."),
        gr.Dropdown(idiomas, label="Idioma origen", value="mul"),
        gr.Dropdown(idiomas, label="Idioma destino", value="en")
    ],
    outputs=gr.Textbox(label="Traducción"),
    title="Traductor Multilenguaje OPUS-MT",
    description="Traduce textos largos entre múltiples idiomas usando OPUS-MT en CPU. Optimizado para uso comercial."
)

iface.launch()