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:
        # Agregar palabras clave médicas a la pregunta
        termino = f"{pregunta} AND (medical OR health OR disease)"
        
        handle = Entrez.esearch(db="pubmed", term=termino, 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[:500]}...\n")  # Aumentar el límite a 500 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:
        # Agregar palabras clave médicas a la pregunta
        termino = f"{pregunta} medical OR healthcare OR treatment"
        url = f"https://api.duckduckgo.com/?q={termino}&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"][:300]  # Aumentar el límite a 300 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:
        # Truncar la pregunta si es demasiado larga
        pregunta = pregunta[:500]  # Limitar a 500 caracteres
        
        inputs = tokenizer(
            f"Explica {pregunta} de forma clara y detallada. Incluye causas, síntomas, diagnóstico y tratamiento si corresponde.",
            return_tensors="pt",
            max_length=512,  # Limitar la longitud máxima de entrada
            truncation=True
        ).to("cpu")
        
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=400,  # Aumentar el límite para respuestas más largas
                do_sample=True,      # Usar muestreo para respuestas más naturales
                temperature=0.5,     # Reducir la creatividad para respuestas más precisas
                num_beams=6          # Mejorar la coherencia
            )
        
        respuesta = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Eliminar el prompt de la respuesta
        respuesta = respuesta.replace(f"Explica {pregunta} de forma clara y detallada. Incluye causas, síntomas, diagnóstico y tratamiento si corresponde.", "").strip()
        
        return respuesta

    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: #f9f9f9;
    font-family: Arial, sans-serif;
}
.gr-button {
    background-color: #007bff;
    color: white;
    border-radius: 5px;
}
"""

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