import cv2
import numpy as np
import gradio as gr
from ultralytics import YOLO
from collections import Counter

def predict(image_path: str):
    # Cargar el modelo
    model = YOLO("best.pt")

    # Leer la imagen desde la ruta
    image = cv2.imread(image_path)

    # Realizar la predicción
    results = model.predict(source=image_path)

    # Lista para almacenar las etiquetas predichas
    predicted_labels = []

    for r in results:
        # Obtener las etiquetas de las predicciones
        labels = r.names  # Nombres de las clases predichas
        classes = r.boxes.cls  # IDs de las clases predichas
        
        # Se guardan las etiquetas de las predicciones
        for cls_id in classes:
            predicted_labels.append(labels[int(cls_id)])

    # Contar las ocurrencias de cada diente en predicted_labels
    predicted_counts = Counter(predicted_labels)

    # Filtrar dientes identificados más de una vez
    identified_multiple = {tooth for tooth, count in predicted_counts.items() if count > 1}

    # Convertir predicted_labels a un conjunto para comparación
    predicted_set = set(predicted_labels)

    # Encontrar dientes que no se identificaron
    missing_teeth = [tooth for tooth in fdi_teeth if tooth not in predicted_set]

    # Ordenar los dientes según el sistema FDI
    def fdi_sort_key(tooth):
        quadrant = int(tooth[0])
        position = int(tooth[1])
        return (quadrant, position)

    # Generar el texto con los resultados
    results_text = []

    results_text.append("Identified Teeth:")
    for tooth in sorted(predicted_set, key=fdi_sort_key):
        results_text.append(f"-{tooth}")

    results_text.append("\nUnidentified  Teeth:")
    for tooth in sorted(missing_teeth, key=fdi_sort_key):
        results_text.append(f"-{tooth}")

    results_text.append("\nTeeth requiring revision:")
    for tooth in sorted(identified_multiple, key=fdi_sort_key):
        results_text.append(f"-{tooth}")



    # Convertir la lista de resultados a una sola cadena de texto
    results_text_str = "\n".join(results_text)

    # Devolver la imagen con las predicciones y el texto de resultados
    # Convertir la imagen con las predicciones a formato numpy para Gradio
    output_image = results[0].plot()  # Usar la primera imagen de resultados
    output_image = cv2.cvtColor(output_image, cv2.COLOR_BGR2RGB)  # Convertir BGR a RGB

    return output_image, results_text_str

# Lista de dientes en el sistema FDI
fdi_teeth = [
    '11', '12', '13', '14', '15', '16', '17', '18',
    '21', '22', '23', '24', '25', '26', '27', '28',
    '31', '32', '33', '34', '35', '36', '37', '38',
    '41', '42', '43', '44', '45', '46', '47', '48'
]

# CSS para el fondo de la interfaz
css = f"""
.gradio-container {{
    background-color: #ADD8E6; /* Fondo azul claro */
}}
"""

# Crear la interfaz de Gradio
with gr.Blocks(css=css) as demo:
    # Encabezado
    gr.Markdown("<h1 style='text-align: center;'>Welcome to DentAIxpert</h1>"
                "<p style='text-align: center;'>Please upload the radiograph you want to analyze.</p>")

    # Contenedor para la carga de imágenes y resultados
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="filepath", label="Input Image")
        with gr.Column():
            output_image = gr.Image(type="numpy", label="Output Image")
            results_text = gr.Textbox(label="Analysis Results", lines=10, placeholder="Results will appear here...")

    analyze_button = gr.Button("Submit")
    analyze_button.click(fn=predict, inputs=input_image, outputs=[output_image, results_text])

demo.launch()