import gradio as gr
import torch
from transformers import pipeline
import time
from functools import wraps
import sys
import multimolecule  # Importar para registrar los modelos de multimolecule
import spaces  # Asegúrate de que este módulo esté correctamente instalado

# Decorador para medir el tiempo de ejecución
def medir_tiempo(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        inicio = time.time()
        resultado = func(*args, **kwargs)
        fin = time.time()
        tiempo_transcurrido = fin - inicio
        print(f"Tiempo de ejecución de '{func.__name__}': {tiempo_transcurrido:.2f} segundos")
        return resultado
    return wrapper

# Decorador para asignar GPU (según la funcionalidad proporcionada por spaces)
@spaces.GPU
@medir_tiempo
def predecir_fill_mask(secuencias):
    """
    Función que realiza una predicción de Fill-Mask para las secuencias de ARN proporcionadas.
    """
    try:
        if not secuencias.strip():
            return "Por favor, ingresa una o más secuencias de ARN válidas con <mask> para predecir."

        # Separar las secuencias por líneas y eliminar espacios vacíos
        secuencias_lista = [seq.strip().upper() for seq in secuencias.strip().split('\n') if seq.strip()]
        resultados = []

        for seq in secuencias_lista:
            # Asegurarse de que la secuencia contenga al menos un <mask>
            if "<MASK>" not in seq and "<mask>" not in seq:
                resultados.append(f"Secuencia sin token <mask>: {seq}. Agrega <mask> donde desees predecir.")
                continue

            # Reemplazar <mask> con el token de máscara del modelo
            seq = seq.replace("<mask>", mask_token).replace("<MASK>", mask_token)

            # Realizar la predicción de Fill-Mask
            try:
                predictions = fill_mask(seq)
            except Exception as e:
                resultados.append(f"Error al predecir para la secuencia: {seq}\n{e}")
                continue

            # Formatear las predicciones
            pred_str = ""
            for pred in predictions:
                pred_str += f"Predicción: {pred['sequence']}, Score: {pred['score']:.4f}\n"

            resultados.append(f"Secuencia: {seq}\n{pred_str}")

        return "\n\n".join(resultados)

    except Exception as e:
        print(f"Error durante la predicción: {e}")
        return f"Error al realizar la predicción: {e}"

# Configurar el dispositivo
device = 0 if torch.cuda.is_available() else -1
if device == -1:
    print("Advertencia: CUDA no está disponible. Se usará la CPU, lo que puede ser lento.")

# Cargar el pipeline de Fill-Mask
try:
    print("Cargando el pipeline de Fill-Mask...")
    fill_mask = pipeline('fill-mask', model='multimolecule/mrnafm', device=device)
except Exception as e:
    print(f"Error al cargar el pipeline de Fill-Mask: {e}")
    sys.exit(1)

# Obtener el token de máscara del modelo
mask_token = fill_mask.tokenizer.mask_token
print(f"Mask token utilizado por el modelo: {mask_token}")

# Definir la interfaz de Gradio
titulo = "OmniGenome: Predicción de Fill-Mask para Secuencias de ARN"
descripcion = (
    "Ingresa una o más secuencias de ARN (una por línea) con un token <mask> donde deseas realizar la predicción. "
    "El modelo utilizado es mRNA-FM de MultiMolecule, un modelo pre-entrenado de lenguaje para secuencias de ARN."
)

iface = gr.Interface(
    fn=predecir_fill_mask,
    inputs=gr.Textbox(
        lines=10, 
        placeholder="Escribe tus secuencias de ARN aquí, una por línea, incluyendo <mask> donde desees predecir...", 
        label="Secuencias de ARN con <mask>"
    ),
    outputs=gr.Textbox(label="Predicciones de Fill-Mask"),
    title=titulo,
    description=descripcion,
    examples=[
        [
            "AUGGCUACUUU<mask>G",
            "GCGCGAU<mask>CGACGUAGCUAGC"
        ],
        [
            "AUAUGCGGUAUCGU<mask>GUA",
            "GGAUACGUGAU<mask>GCUAGCAGU"
        ]
    ],
    cache_examples=False,
    allow_flagging="never"
)

# Ejecutar la interfaz
if __name__ == "__main__":
    iface.launch()