import gradio as gr
import cv2
from matplotlib import pyplot as plt
import numpy as np
from typing import List
from ultralytics import YOLO
import cv2
import numpy as np
import os
import shutil


def recortar_notas(imagen_path: str) -> int:
    if os.path.exists('deteccion'):
        shutil.rmtree('deteccion')
    if os.path.exists('diarios'):
        shutil.rmtree('diarios')
    if os.path.exists('recorte'):
        shutil.rmtree('recorte')
    
    model = YOLO("seg-640.pt")
    results = model.predict(source=imagen_path, save=True, save_txt=True, project="diarios", classes=0, conf=0.5) 

    # cargar imagen tif
    image = cv2.imread(imagen_path)

    # crear una carpeta para guardar los recortes
    filename = os.path.splitext(os.path.basename(results[0].path))[0]
    foldername = os.path.splitext(filename)[0]
    if not os.path.exists(f"recorte/{foldername}"):
        os.makedirs(f"recorte/{foldername}")

    datos = results[0].masks.xy

    num_notas = 0

    for i, coords in enumerate(datos): 
        # coordenadas de segmentación
        coordenadas = np.array(coords)

        # crear una máscara vacía
        mask = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8)

        # dibujar un polígono en la máscara
        cv2.fillPoly(mask, [coordenadas.astype(np.int32)], 255)

        # aplicar la máscara a la imagen
        masked = cv2.bitwise_and(image, image, mask=mask)

        # recortar la imagen utilizando la caja delimitadora de la máscara
        x, y, w, h = cv2.boundingRect(mask)
        recorte = masked[y:y+h, x:x+w]

        # guardar el recorte como una imagen
        recorte_path = os.path.join(f"recorte/{foldername}", f"nota {i}.jpg")
        cv2.imwrite(recorte_path, recorte)

        num_notas += 1
    
    noticias(foldername)
    segmentacion = f"diarios/predict/{filename}.png"
    segmentacion = cv2.imread(segmentacion)
    recorte1 = f"recorte/{foldername}/nota 0.jpg"
    prediccion = f"deteccion/{foldername}/predict/nota 0.jpg"
        
    try:
        recorte2 = f"recorte/{foldername}/nota 1.jpg"
        prediccion2 = f"deteccion/{foldername}/predict/nota 1.jpg"
        if not os.path.exists(recorte2):
            raise Exception(f"El archivo {recorte2} no existe.")
    except:
          recorte2 = "sin_nota.jpg"
          prediccion2 = "sin_nota.jpg"
    
    return segmentacion, prediccion, prediccion2, num_notas

def noticias(carpeta):
        nombre = carpeta
        model = YOLO("best-detect.pt")
        results = model.predict(source=f"recorte/{nombre}", save=True, save_txt=True,project=f"deteccion/{nombre}", conf=0.75)  # save plotted images
        return print(f"Imagen {nombre} procesada correctamente")

with gr.Blocks() as demo:
    gr.Markdown(
        """
        <h1 align="center"> IA por la Identidad
        </h1>
        <h2 align="center"> Dathaton 
        </h2>
        """
    )   
    gr.Markdown(
        """
<p align="center">
    <img width = 600 src="https://raw.githubusercontent.com/BonfantiMatias/images/main/banner%20fundaciones.jpeg">
    
</p>

"""
    )

    
    gr.Markdown(
    """
- Puede Seleccionar una de las imagenes de ejemplo o subir una desde su pc
- Para borrar la imagen que esta en la ventana de procesamiento debe presionar la "X" que se encuentra en el vertice superior derecho
"""
    )
    with gr.Row():
            seed = gr.components.Image(type="filepath", label="Input") 
            with gr.Row():
                with gr.Column():
                        gr.Examples(["Ambito2.jpg"], inputs=[seed])
                        gr.Examples(["Clarin2.jpg"], inputs=[seed])
                        
                with gr.Column():
                        gr.Examples(["Popular.jpg"], inputs=[seed])
                        gr.Examples(["Clarin3.jpg"], inputs=[seed])

    with gr.Row():
                notas = gr.Label(label="Numero de Notas")
      
    with gr.Row():
            
        with gr.Column():
            segmentacion = gr.Image(label="Segmentacion Notas")
        with gr.Column():
            prediccion = gr.Image(label="Prediccion Primera Nota")
        with gr.Column():
            prediccion2 = gr.Image(label="Prediccion Segunda Nota")

    with gr.Row():
        btn = gr.Button("Procesar Imagen")
        btn.click(recortar_notas, inputs=[seed], outputs=[segmentacion, prediccion, prediccion2, notas])          
          
if __name__ == "__main__":
    demo.launch()