import gradio as gr
from setfit import SetFitModel
# Import libraries
import PyPDF2
import openpyxl
import os
import glob
import re
import itertools
import platform
from tempfile import TemporaryDirectory
from pathlib import Path
import pytesseract
from pdf2image import convert_from_path
from PIL import Image

#print('HOLA:', os.walk('../'))

def ocrtotext(filename):
    # Almacena todas las páginas del PDF en una variable
    image_file_list = []  # Lista para almacenar los nombres de los archivos de imagen
    PDF_file = filename  # Nombre del archivo PDF
    text_file = (f'{filename}.txt')  # Nombre del archivo de texto de salida

    # Parte #1: Convirtiendo el PDF a imágenes
    pdf_pages = convert_from_path(PDF_file, 300)
    # Lee el archivo PDF a 300 DPI

    # Itera a través de todas las páginas almacenadas arriba
    for page_enumeration, page in enumerate(pdf_pages, start=1):
        # Crea un nombre de archivo para almacenar la imagen
        filename = f"page_{page_enumeration:03}.jpg"

        # Guarda la imagen de la página en el sistema
        page.save(filename, "JPEG")
        image_file_list.append(filename)

    # Parte #2: Reconociendo texto desde las imágenes usando OCR
    output = ''  # Variable para almacenar el texto reconocido

    # Abre el archivo en modo de apendizaje para agregar contenido de todas las imágenes al mismo archivo
    with open(text_file, "a") as output_file:

        # Itera desde 1 hasta el número total de páginas
        for image_file in image_file_list:

            # Reconoce el texto en la imagen usando pytesseract
            #text = str(((pytesseract.image_to_string(Image.open(image_file), lang="spa"))))
            text = str(((pytesseract.image_to_string(Image.open(image_file)))))

            # Elimina los guiones al final de las líneas
            text = text.replace("-\n", "")

            # Agrega el texto reconocido a la variable de salida
            output += text

        # Escribe el texto reconocido en el archivo de salida
        output_file.write(output)
        return output

def cortar_en_bloques(texto, longitud_bloque):
    palabras = texto.split()
    bloques = []
    bloque_actual = []

    for palabra in palabras:
        bloque_actual.append(palabra)

        if len(bloque_actual) == longitud_bloque:
            bloques.append(" ".join(bloque_actual))
            bloque_actual = []

    # Si queda un bloque parcial al final, agregarlo
    if bloque_actual:
        bloques.append(" ".join(bloque_actual))

    return bloques

# Cargar el modelo
model = SetFitModel.from_pretrained("desarrolloasesoreslocales/SetFitPruebaRecorte")

# Mapear las etiquetas

# Definir la función de predicción
def predict(file):
    
  recorte_general = ""
  ocr_text = ocrtotext(file.name)
  # Crear chunks
  chunks = cortar_en_bloques(ocr_text, 150)
  first = -1
  margin = int(len(chunks) * 0.25)
  chunks_removable = chunks[:margin] + chunks[-margin:]

  for i in range(len(chunks)):
    print('Recortando -', round((i/len(chunks))*100), '%')
    if chunks[i] not in chunks_removable or model.predict([chunks[i]]).item() == 1:
      if first == -1:
        first = i
      recorte_general += chunks[i] + " "

  if first > 0:
    recorte_general = chunks[first-1] + recorte_general
  print(100, '%')

  recorte_final = ""

  # Definir tamñano de fragmentos de texto
  # text_splitter2 = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=0, length_function=len)
  # Crear chunks
  chunks2 = cortar_en_bloques(recorte_general, 80)
  margin_s = int(len(chunks2) * 0.1)
  margin_e = int(len(chunks2) * 0.1)
  # if margin_s > 1:
  chunks_removable2 = chunks2[:margin_s] + chunks2[-margin_e:]
  # chunks_removable2 = chunks2[-margin_e:]

  for i in range(len(chunks2)):
    print('Recortando -', round((i/len(chunks2))*100), '%')
    if chunks2[i] not in chunks_removable2 or model.predict([chunks2[i]]).item() == 1:
      recorte_final += chunks2[i] + " "
  print(100, '%')
  return recorte_final

# Crear una interfaz Gradio
iface = gr.Interface(
    fn=predict, 
    inputs=gr.inputs.File(), 
    outputs=gr.Textbox(), 
    live=False,
    title="Recortador de Texto"
)

# Iniciar la interfaz Gradio
iface.launch()