from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import gradio as gr
from Bio import Entrez
import requests
from functools import lru_cache
import logging

# Configurar logging
logging.basicConfig(level=logging.INFO)

# Cargar un modelo de lenguaje médico más pequeño (optimizado para CPU)
model_name = "microsoft/BioGPT"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name).to("cpu")  # Usar CPU

# Configurar PubMed
Entrez.email = "tu_correo@example.com"  # Proporciona un correo válido

# Función para buscar en PubMed con caché
@lru_cache(maxsize=100)
def buscar_en_pubmed(pregunta):
    try:
        handle = Entrez.esearch(db="pubmed", term=pregunta, retmax=1)  # Buscar solo 1 artículo
        record = Entrez.read(handle)
        handle.close()
        
        if record["IdList"]:
            referencias = []
            for id_articulo in record["IdList"]:
                handle = Entrez.efetch(db="pubmed", id=id_articulo, rettype="abstract", retmode="text")
                resumen = handle.read()
                referencias.append(f"Artículo {id_articulo}:\n{resumen[:300]}...\n")  # Limitar a 300 caracteres
                handle.close()
            return "\n".join(referencias)
        else:
            return "No encontré artículos relevantes en PubMed."
    except Exception as e:
        logging.error(f"Error al buscar en PubMed: {e}")
        return "Error al buscar en PubMed. Inténtalo de nuevo."

# Función para buscar en Internet con tiempo de espera
def buscar_en_internet(pregunta):
    try:
        url = f"https://api.duckduckgo.com/?q={pregunta}&format=json"
        respuesta = requests.get(url, timeout=10)  # Incrementa el tiempo de espera a 10 segundos
        datos = respuesta.json()
        
        if "Abstract" in datos and datos["Abstract"]:
            return datos["Abstract"][:200]  # Limitar a 200 caracteres
        else:
            return "No encontré información en Internet."
    except Exception as e:
        logging.error(f"Error al buscar en Internet: {e}")
        return "Error al buscar en Internet. Inténtalo de nuevo."

# Función para generar respuestas del modelo (optimizada para CPU)
def generar_respuesta(pregunta):
    try:
        inputs = tokenizer(
            f"Eres un profesor médico con vasto conocimiento en fisiología, bioquímica, farmacología y otras áreas médicas. "
            f"Explica de manera clara, sencilla y didáctica. Utiliza términos médicos y explícalos dentro de tus respuestas. "
            f"Si no estás seguro de algo, di 'No tengo suficiente información para responder a eso'. Pregunta: {pregunta}",
            return_tensors="pt"
        ).to("cpu")
        
        with torch.no_grad():
            outputs = model.generate(**inputs, max_length=100)  # Reducir max_length para mayor velocidad
        
        return tokenizer.decode(outputs[0], skip_special_tokens=True)

    except Exception as e:
        logging.error(f"Error al generar respuesta del modelo: {e}")
        return "No pude generar una respuesta. Inténtalo de nuevo."

# Función principal para hacer preguntas al bot
def pregunta_medica_con_carga(pregunta):
    try:
        # Validar la entrada
        if not isinstance(pregunta, str) or not pregunta.strip():
            yield "Por favor, ingresa una pregunta válida."
            return
        
        yield "Procesando tu pregunta..."  # Mensaje de carga
        
        # Respuesta del modelo médico
        respuesta_modelo = generar_respuesta(pregunta)
        
        # Buscar en PubMed
        referencias_pubmed = buscar_en_pubmed(pregunta)
        
        # Buscar en Internet
        respuesta_internet = buscar_en_internet(pregunta)
        
        # Combinar respuestas
        respuesta_final = (
            f"Respuesta del Profesor Médico:\n{respuesta_modelo}\n\n"
            f"Referencias de PubMed:\n{referencias_pubmed}\n\n"
            f"Información de Internet:\n{respuesta_internet}"
        )
        
        yield respuesta_final

    except Exception as e:
        logging.error(f"Error inesperado: {e}")
        yield "Lo siento, ocurrió un error al procesar tu pregunta. Por favor, inténtalo de nuevo más tarde."

# Crear la interfaz web con Gradio
css = """
.gradio-container {
    background-color: #f0f8ff;
}
"""

interfaz = gr.Interface(
    fn=pregunta_medica_con_carga,
    inputs=gr.Textbox(label="Pregunta"),
    outputs=gr.Textbox(label="Respuesta"),
    title="Profesor Médico Bot",
    description="Pregúntale al Profesor Médico sobre medicina. Te explicará de manera clara y didáctica.",
    css=css
)

# Lanzar la interfaz
interfaz.launch()