import gradio as gr
import torch
from transformers import pipeline
import time
from functools import wraps
import sys
from multimolecule import RnaTokenizer, RnaFmModel  # Importar clases específicas de multimolecule

# 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

# Verificar si CUDA está disponible para el modelo principal
device = "cuda" if torch.cuda.is_available() else "cpu"
if device == "cpu":
    print("Advertencia: CUDA no está disponible. Se usará la CPU, lo que puede ser lento.")

# Cargar el modelo y el tokenizador
model_name = "multimolecule/mrnafm"

try:
    print("Cargando el tokenizador...")
    tokenizer = RnaTokenizer.from_pretrained(model_name)
except ValueError as e:
    print(f"Error al cargar el tokenizador: {e}")
    sys.exit(1)

try:
    print("Cargando el modelo...")
    model = RnaFmModel.from_pretrained(model_name)
    model.to(device)
except Exception as e:
    print(f"Error al cargar el modelo: {e}")
    sys.exit(1)

@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 = []

        # Crear el pipeline de fill-mask utilizando el tokenizador y modelo cargados
        fill_mask = pipeline('fill-mask', model=model, tokenizer=tokenizer, device=0 if device == "cuda" else -1)

        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

            # Realizar la predicción de fill-mask
            predictions = fill_mask(seq)

            # 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}"

# 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()