import gradio as gr
import os
import fitz  # PyMuPDF
from groq import Groq
from langchain_groq import ChatGroq
import json
import logging

# Configuración de logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

RAID_WEIGHTS = {
    "Dolor": 0.21,
    "Discapacidad Funcional": 0.16,
    "Fatiga": 0.15,
    "Sueño": 0.12,
    "Bienestar Físico": 0.12,
    "Bienestar Emocional": 0.12,
    "Afronte": 0.12
}

RAID_PROMPT_TEMPLATE = """
Analiza la siguiente conversación y registro clínico del paciente con artritis reumatoide.

Si tanto el "Texto del paciente" como el "Registro clínico" están vacíos o no contienen información relevante, responde únicamente con el siguiente mensaje:
    "No hay información suficiente para calcular el RAID."

En caso de haber información, extrae y asigna una puntuación del 0 al 10 para cada una de las siguientes categorías, basándote en los detalles proporcionados:

1. Dolor:
   - Circule el número que mejor describa el dolor que sintió debido a su RA durante la última semana.
   - Escala: 0 = Ninguno, 10 = Extremo.
2. Discapacidad Funcional:
   - Circule el número que mejor describa la dificultad que tuvo para realizar actividades diarias debido a su RA durante la última semana.
   - Escala: 0 = Sin dificultad, 10 = Dificultad extrema.
3. Fatiga:
   - Circule el número que mejor describa la fatiga que experimentó debido a su RA durante la última semana.
   - Escala: 0 = Sin fatiga, 10 = Totalmente exhausto.
4. Sueño:
   - Circule el número que mejor describa las dificultades para dormir que experimentó durante la última semana.
   - Escala: 0 = Sin dificultad, 10 = Dificultad extrema.
5. Bienestar Físico:
   - Circule el número que mejor describa su nivel de bienestar físico en relación a su RA durante la última semana.
   - Escala: 0 = Muy bueno, 10 = Muy malo.
6. Bienestar Emocional:
   - Circule el número que mejor describa su nivel de bienestar emocional en relación a su RA durante la última semana.
   - Escala: 0 = Muy bueno, 10 = Muy malo.
7. Afronte:
   - Circule el número que mejor describa cómo enfrentó o se adaptó a su enfermedad durante la última semana.
   - Escala: 0 = Muy bien, 10 = Muy mal.

Utiliza los siguientes pesos para calcular el valor final del RAID:
{weights}


La fórmula de cálculo es:
    RAID final = (Dolor × 0.21) + (Discapacidad Funcional × 0.16) + (Fatiga × 0.15) + (Bienestar Físico × 0.12) + (Sueño × 0.12) + (Bienestar Emocional × 0.12) + (Afronte × 0.12)

El valor final estará en el rango de 0 a 10, donde valores más altos indican un peor estado.

Por favor, responde únicamente con un JSON válido que incluya:
    - "raid_scores": un objeto que contenga las puntuaciones asignadas para cada categoría.
    - "raid_total": el valor final calculado.
Si falta información para alguna categoría, asigna 0 a esa categoría y añade una clave "nota" con un mensaje indicando que la información es parcial.

Ejemplo de respuesta JSON:
{{
    "raid_scores": {{
        "Dolor": 8,
        "Discapacidad Funcional": 6,
        "Fatiga": 7,
        "Sueño": 5,
        "Bienestar Físico": 6,
        "Bienestar Emocional": 7,
        "Afronte": 4
    }},
    "raid_total": 6.45
}}

Texto del paciente:
{patient_text}

Registro clínico:
{clinical_record}
"""


def evaluate_raid(patient_text, clinical_record):
    """Evalúa el RAID score basado en el texto del paciente y registro clínico.
    Si la respuesta es parcial, se asigna 0 a las categorías faltantes y se añade
    una clave 'error' con la descripción del problema."""
    try:
        prompt = RAID_PROMPT_TEMPLATE.format(
            patient_text=patient_text[:2000],  # Limitar texto para contexto
            clinical_record=clinical_record[:2000],
            weights=json.dumps(RAID_WEIGHTS, ensure_ascii=False)
        )
        logger.info("Evaluando RAID con prompt: %s", prompt)
        
        response = chat_groq.invoke(prompt)
        content = response.content
        logger.info("Respuesta de chat_groq: %s", content)
        
        # Extraer JSON de la respuesta
        start = content.find('{')
        end = content.rfind('}') + 1
        json_str_candidate = content[start:end]
        try:
            json_response = json.loads(json_str_candidate)
        except json.JSONDecodeError:
            # Intentar extraer contenido entre marcadores ```json
            try:
                json_str_candidate = content.split('```json')[1].split('```')[0].strip()
                json_response = json.loads(json_str_candidate)
            except Exception as e:
                logger.error("Error al extraer JSON usando marcadores: %s", e)
                json_response = {}
        
        if not isinstance(json_response, dict):
            logger.warning("Respuesta JSON no es un diccionario. Se usará un diccionario vacío.")
            json_response = {}
        
        # Obtener o inicializar raid_scores
        raid_scores = json_response.get('raid_scores', {})
        if not isinstance(raid_scores, dict):
            logger.warning("El campo 'raid_scores' no es un diccionario, se establecerá como vacío.")
            raid_scores = {}
        
        # Completar las categorías faltantes con 0
        for category in RAID_WEIGHTS:
            if category not in raid_scores:
                logger.warning("Categoría '%s' no encontrada en la respuesta. Se asignará 0.", category)
                raid_scores[category] = 0
        
        # Calcular el puntaje total usando las categorías disponibles
        calculated_total = sum(raid_scores[k] * RAID_WEIGHTS[k] for k in RAID_WEIGHTS)
        raid_total = round(calculated_total, 2)
        
        # Actualizar json_response
        json_response['raid_scores'] = raid_scores
        json_response['raid_total'] = raid_total
            
        return json_response
    
    except Exception as e:
        logger.error("Error en evaluación RAID: %s", e, exc_info=True)
        # Retorna resultados parciales con mensaje de error
        partial_result = {
            "raid_scores": {k: 0 for k in RAID_WEIGHTS},
            "raid_total": 0,
            "error": str(e)
        }
        return partial_result

def format_raid_results(raid_data):
    """Formatea los resultados del RAID para visualización.
    Si se produjo un error, se añade una nota indicando el problema."""
    if not raid_data:
        return "No se pudo calcular el RAID score"
    
    scores = raid_data.get('raid_scores', {})
    total = raid_data.get('raid_total', 0)
    
    breakdown_lines = []
    for category, weight in RAID_WEIGHTS.items():
        score = scores.get(category, 0)
        breakdown_lines.append(f"- {category}: {score} (Peso: {weight*100}%)")
    breakdown = "\n".join(breakdown_lines)
    
    interpretation = "Interpretación del puntaje:\n"
    if total >= 7:
        interpretation += "RAID alto: Impacto significativo en la calidad de vida. Considerar ajuste terapéutico."
    elif total >= 4:
        interpretation += "RAID moderado: Impacto notable. Monitoreo cercano recomendado."
    else:
        interpretation += "RAID bajo: Buen control sintomático. Mantener seguimiento."
    
    result = f"""
**Puntaje Total RAID: {total:.2f}/10**

**Desglose:**
{breakdown}

{interpretation}
"""
    if "error" in raid_data:
        result += f"\n\n**Nota:** Se produjo un error durante el cálculo del RAID: {raid_data['error']}"
    return result

# Inicialización del cliente
api_key = os.environ.get("GROQ_API_KEY")
if not api_key:
    raise ValueError("GROQ_API_KEY no está configurada en las variables de entorno.")
client = Groq(api_key=api_key)
model_name = "llama-3.3-70b-versatile"
chat_groq = ChatGroq(model=model_name)

# Funciones de procesamiento
def transcribe_audio(audio_filepath):
    if not audio_filepath:
        return ""
    try:
        with open(audio_filepath, "rb") as file:
            transcription = client.audio.transcriptions.create(
                file=(audio_filepath, file.read()),
                model="whisper-large-v3",
                response_format="json",
                temperature=0.0
            )
        logger.info("Transcripción de audio exitosa.")
        return transcription.text
    except Exception as e:
        logger.error("Error en transcripción de audio: %s", e, exc_info=True)
        return ""

def extract_text_from_pdf(pdf_path):
    try:
        text = ''
        with fitz.open(pdf_path) as doc:
            for page in doc:
                text += page.get_text()
        logger.info("Extracción de texto de PDF exitosa para: %s", pdf_path)
        return text
    except Exception as e:
        logger.error("Error al extraer texto de PDF %s: %s", pdf_path, e, exc_info=True)
        return ""

def extract_texts_from_pdfs(pdfs):
    text = ''
    if not pdfs:
        return text
    for pdf in pdfs:
        if isinstance(pdf, dict) and 'name' in pdf:
            pdf_path = pdf['name']
        elif isinstance(pdf, str):
            pdf_path = pdf
        else:
            continue
        pdf_text = extract_text_from_pdf(pdf_path)
        text += pdf_text + "\n"
    return text

def split_text_into_chunks(text, max_words_per_chunk):
    words = text.split()
    chunks = []
    for i in range(0, len(words), max_words_per_chunk):
        chunk = ' '.join(words[i:i + max_words_per_chunk])
        chunks.append(chunk)
    return chunks

def organize_clinical_record(current_text, transcription_text, pdf_text):
    clinical_record_template = """
MOTIVO DE CONSULTA: usa una frase en palabras del paciente entre comillas
ENFERMEDAD ACTUAL:
(usa terminología médica. En orden cronológico desde el inicio de los síntomas, no incluir la edad ni los antecedentes en esta sección, evolución de los síntomas, factores desencadenantes, hitos de la enfermedad del paciente, finaliza describiendo cómo se siente hoy)
REVISIÓN POR SISTEMAS:
(usa terminología médica)
ANTECEDENTES:
**Patológicos: (describir en lenguaje técnico médico la enfermedad con clasificación y complicaciones relacionadas de cada antecedente)
**Alérgicos: (con tipo de reacción y a cuál medicamento)
**Tóxicos: (exposición a biomasa, IPA, Alcohol por unidades estándar, otros, naturales)
**Familiares:
**Transfusionales:
**Traumáticos:
**Ginecológicos:
**Quirúrgicos: (fecha y procedimiento)
**Estado de vacunación:
**Hospitalizaciones previas (fecha y descripción breve)
**Medicamentos:
AYUDAS DIAGNÓSTICAS:
(ordenar todas las ayudas diagnósticas por fecha de forma que sea simple y sencillo leer los resultados para el médico, cuando se requiera presenta los resultados en miles ya x10e3 o similares; es decir, multiplica el resultado por 1000, asegurándote de no omitir ninguna ayuda, y no interpretes, solo pon los valores sin rango de normalidad en prosa. Separa cada examen con una coma, usa minúsculas y organiza por fechas.
Por ejemplo: 
11/10/2024: resultado 1 , resultado 2, ...
12/11/2023: resultado 1 , resultado 2, ...)
"""
    prompt = f"""
Toma el siguiente borrador del registro clínico y actualízalo con la nueva información proporcionada, siguiendo la estructura dada:
Estructura del Registro Clínico:
{clinical_record_template}
Borrador Actual del Registro Clínico:
{current_text}
Nueva Información de Audio:
{transcription_text}
Nueva Información del PDF:
{pdf_text}
Actualiza el borrador incorporando la nueva información en las secciones correspondientes, sin eliminar información previa que aún sea relevante.
"""
    try:
        organized_text = chat_groq.invoke(prompt)
        logger.info("Organización del registro clínico exitosa.")
        return organized_text.content
    except Exception as e:
        logger.error("Error al invocar ChatGroq para organizar el registro clínico: %s", e, exc_info=True)
        return current_text  # Se retorna el texto actual si falla la invocación

def process_input(audio, pdfs, current_text):
    try:
        transcription_text = transcribe_audio(audio)
    except Exception as e:
        transcription_text = ""
        logger.error("Error en transcripción de audio: %s", e, exc_info=True)
    
    try:
        pdf_text = extract_texts_from_pdfs(pdfs)
    except Exception as e:
        pdf_text = ""
        logger.error("Error en extracción de PDFs: %s", e, exc_info=True)

    updated_text = current_text
    raid_results = None

    # Procesamiento en lotes
    max_chunk_words = 2500
    
    for text_label, text_content in [("Audio", transcription_text), ("PDF", pdf_text)]:
        if not text_content:
            continue

        text_chunks = split_text_into_chunks(text_content, max_chunk_words)

        for chunk in text_chunks:
            transcription_chunk = chunk if text_label == "Audio" else ""
            pdf_chunk = chunk if text_label == "PDF" else ""
            
            # Actualizar registro clínico
            organized_record = organize_clinical_record(updated_text, transcription_chunk, pdf_chunk)
            if organized_record:
                updated_text = organized_record
                
            # Evaluar RAID solo con el audio (conversación actual)
            if text_label == "Audio" and chunk:
                raid_results = evaluate_raid(chunk, updated_text)

    return updated_text, format_raid_results(raid_results), ""

# Configuración del tema
theme = gr.themes.Base(
    primary_hue=gr.themes.Color(
        c100="#fce7f3", c200="#fbcfe8", c300="#f9a8d4",
        c400="#f472b6", c50="#fdf2f8", c500="#ff4da6",
        c600="#db2777", c700="#ff299b", c800="#f745b6",
        c900="#ff38ac", c950="#e1377e"
    ),
    secondary_hue="pink",
    neutral_hue="neutral",
)

# Texto inicial
initial_text = """
MOTIVO DE CONSULTA: 

ENFERMEDAD ACTUAL:

REVISIÓN POR SISTEMAS:

ANTECEDENTES:
**Patológicos: 
**Alérgicos: 
**Tóxicos: 
**Familiares:
**Transfusionales:
**Traumáticos: 
**Ginecológicos:
**Quirúrgicos:
**Estado de vacunación: 
**Hospitalizaciones previas:
**Medicamentos:

AYUDAS DIAGNÓSTICAS:
"""

# Interfaz de Gradio
with gr.Blocks(theme=theme) as iface:
    gr.Markdown("# Sistema de Evaluación Reumatológica")
    
    with gr.Row():
        iterative_output = gr.Textbox(
            label="Registro Clínico",
            value=initial_text,
            lines=20,
            elem_id="clinical_record"
        )
        
        raid_output = gr.Markdown(
            "### Resultados RAID\nEsperando evaluación...",
            elem_id="raid_results"
        )
    
    current_state = gr.State(value=initial_text)
    
    with gr.Row():
        audio_filepath = gr.Audio(sources=["microphone"], type="filepath", 
                                  label="Conversación con el Paciente")
        pdf_files = gr.File(file_types=[".pdf"], label="Documentos Médicos", 
                            file_count="multiple")
    
    debug_output = gr.Textbox(label="Registros", lines=5, visible=True)
    
    # Actualizar current_state cuando se edite el registro clínico
    def on_text_change(updated_text):
        return updated_text
    
    iterative_output.change(
        fn=on_text_change,
        inputs=iterative_output,
        outputs=current_state
    )
    
    def on_audio_change(audio_filepath, pdfs, current_text):
        logger.info("on_audio_change: audio_filepath = %s", audio_filepath)
        if not audio_filepath:
            return current_text, "No se proporcionó audio.", current_text
        if not current_text or not current_text.strip():
            current_text = initial_text
        updated_text, debug_info, _ = process_input(audio_filepath, pdfs, current_text)
        return updated_text, debug_info, updated_text
    
    audio_filepath.change(
        fn=on_audio_change,
        inputs=[audio_filepath, pdf_files, iterative_output],
        outputs=[iterative_output, debug_output, current_state]
    )
    
    def on_pdfs_change(audio_filepath, pdfs, current_text):
        logger.info("on_pdfs_change: pdfs = %s", pdfs)
        if not pdfs:
            return current_text, "No se proporcionaron PDFs.", current_text
        if not current_text or not current_text.strip():
            current_text = initial_text
        updated_text, debug_info, _ = process_input(audio_filepath, pdfs, current_text)
        return updated_text, debug_info, updated_text
    
    pdf_files.upload(
        fn=on_pdfs_change,
        inputs=[audio_filepath, pdf_files, iterative_output],
        outputs=[iterative_output, debug_output, current_state]
    )

iface.launch(auth=[("her", "her")])