import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModel
import time
from functools import wraps
import sys
import spaces  # Asegúrate de que este módulo esté disponible y 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

# 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 = "yangheng/OmniGenome"

try:
    print("Cargando el tokenizador...")
    tokenizer = AutoTokenizer.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 = AutoModel.from_pretrained(model_name)
    model.to(device)
except Exception as e:
    print(f"Error al cargar el modelo: {e}")
    sys.exit(1)

@spaces.GPU(duration=120)  # Decorador para asignar GPU durante 120 segundos
@medir_tiempo
def predecir_estructura_rna(secuencias):
    """
    Función que predice estructuras secundarias de ARN a partir de secuencias de ARN proporcionadas.
    """
    try:
        if not secuencias.strip():
            return "Por favor, ingresa una o más secuencias de ARN válidas."

        # 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:
            # Validar la secuencia de ARN
            if not all(residue in 'AUCG' for residue in seq):
                resultados.append(f"Secuencia inválida: {seq}. Solo se permiten los nucleótidos A, U, C y G.")
                continue

            # Tokenizar la secuencia
            inputs = tokenizer(seq, return_tensors="pt")
            input_ids = inputs["input_ids"].to(device)
            attention_mask = inputs["attention_mask"].to(device)

            # Aplicar el modelo para obtener los logits
            with torch.no_grad():
                outputs = model(input_ids=input_ids, attention_mask=attention_mask)

            # Asumimos que el modelo devuelve logits para cada nucleótido que indican la estructura secundaria
            # Debes ajustar esto según la arquitectura específica de OmniGenome

            # Por ejemplo, supongamos que el modelo tiene una cabeza de clasificación con N etiquetas
            # donde cada etiqueta representa una clase de estructura secundaria (e.g., Helix, Loop, etc.)

            # Obtener las predicciones seleccionando la clase con el logit más alto
            predictions = torch.argmax(outputs.logits, dim=-1).squeeze().tolist()

            # Definir el mapeo de clases según la documentación del modelo OmniGenome
            # Este mapeo debe ajustarse a las clases específicas que OmniGenome predice
            # Por ejemplo:
            structure_mapping = {
                0: 'Helix',
                1: 'Loop',
                2: 'Bulge',
                3: 'Internal Loop',
                # Agrega más clases si es necesario
            }

            # Convertir las predicciones numéricas a etiquetas legibles
            predicted_structures = [structure_mapping.get(pred, "Unknown") for pred in predictions]

            # Emparejar cada nucleótido con su etiqueta de estructura predicha
            nucleotide_to_structure = list(zip(list(seq), predicted_structures))

            # Formatear el resultado para mostrarlo en la interfaz
            secuencia_resultado = []
            for i, (nucleotide, structure) in enumerate(nucleotide_to_structure):
                secuencia_resultado.append(f"Posición {i+1} - {nucleotide}: {structure}")

            # Unir las predicciones en un solo string
            resultados.append("\n".join(secuencia_resultado))

        # Unir los resultados de todas las secuencias separadas por dos saltos de línea
        return "\n\n".join(resultados)

    except Exception as e:
        print(f"Error durante la predicción: {e}")
        return f"Error al predecir las estructuras de ARN: {e}"

# Definir la interfaz de Gradio
titulo = "OmniGenome: Predicción de Estructuras Secundarias de ARN"
descripcion = (
    "Ingresa una o más secuencias de ARN (una por línea) y obtén predicciones de estructuras secundarias para cada nucleótido."
    " El modelo utilizado es OmniGenome, un modelo de fundamentos basado en transformadores para alineación secuencia-estructura en tareas genómicas."
)

iface = gr.Interface(
    fn=predecir_estructura_rna,
    inputs=gr.Textbox(
        lines=10, 
        placeholder="Escribe tus secuencias de ARN aquí, una por línea (solo A, U, C, G)...", 
        label="Secuencias de ARN"
    ),
    outputs=gr.Textbox(label="Predicciones de Estructuras Secundarias de ARN"),
    title=titulo,
    description=descripcion,
    examples=[
        [
            "AUGGCUACUUUCG",
            "GCGCGAUCGACGUAGCUAGC"
        ],
        [
            "AUAUGCGGUAUCGUACGUA",
            "GGAUACGUGAUCGUAGCAGU"
        ]
    ],
    cache_examples=False,
    allow_flagging="never"
)

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