import gradio as gr
import spacy
import pandas as pd
from docx import Document
from io import BytesIO
import tempfile
import os
import multiprocessing as mp  # Importar multiprocessing para paralelismo

# Cargar el modelo de SpaCy en español
nlp = spacy.load('zh_core_web_trf')
#nlp.max_length = 15000000  # Aumenta el límite a 3 millones de caracteres

# Función para procesar bloques de texto y extraer nombres de personas
def extract_names_from_text(text):
    print(f'{len(text)}/n/n{text}')
    doc = nlp(text)
    # Extraer las entidades de tipo PERSON
    persons = [ent.text for ent in doc.ents if ent.label_ == 'PERSON']
    return persons

# Función para dividir el texto en fragmentos más pequeños
def split_text(text, max_length=100000):
    result = []
    current_chunk = []
    current_length = 0

    # Divide por salto de línea simple en lugar de doble
    paragraphs = text.split('\n')  # Usamos '\n' ya que en chino no se usan saltos dobles

    for paragraph in paragraphs:
        paragraph_length = len(paragraph) + 1  # Considera el '\n' añadido entre párrafos
        if current_length + paragraph_length <= max_length:
            current_chunk.append(paragraph)
            current_length += paragraph_length
        else:
            # Guarda el fragmento actual y empieza uno nuevo
            result.append('\n'.join(current_chunk))
            current_chunk = [paragraph]
            current_length = paragraph_length

    # Añadir el último fragmento si no está vacío
    if current_chunk:
        result.append('\n'.join(current_chunk))

    return result

# Función para paralelizar la extracción de nombres
def extract_names_from_fragments(fragments):
    # Utiliza todos los núcleos disponibles de la CPU
    with mp.Pool(processes=4) as pool:
        results = pool.map(extract_names_from_text, fragments)
    return results

# Función principal para extraer nombres de personas desde un archivo DOCX
def extract_names_from_docx(docx_file):
    # Cargar el archivo DOCX
    document = Document(docx_file)
    full_text = []
    for para in document.paragraphs:
        full_text.append(para.text)
    
    # Unir todo el texto
    text = ' '.join(full_text)
    
    # Dividir el texto en fragmentos si es necesario
    text_fragments = split_text(text)
    
    # Extraer los nombres de cada fragmento en paralelo
    results = extract_names_from_fragments(text_fragments)
    
    # Unir todos los resultados de nombres en una sola lista
    all_persons = []
    for persons in results:
        all_persons.extend(persons)
    
    # Eliminar duplicados
    all_persons = list(set(all_persons))
    
    # Crear un DataFrame
    df = pd.DataFrame(all_persons, columns=['Nombres'])
    
    # Crear un archivo temporal para guardar el Excel
    temp_dir = tempfile.mkdtemp()
    temp_file_path = os.path.join(temp_dir, "nombres_personas.xlsx")
    
    # Guardar el DataFrame en un archivo Excel
    with pd.ExcelWriter(temp_file_path, engine='openpyxl') as writer:
        df.to_excel(writer, index=False)
    
    return temp_file_path  # Devolver la ruta del archivo temporal

# Interfaz de Gradio
iface = gr.Interface(
    fn=extract_names_from_docx,
    inputs=gr.File(file_types=[".docx"]),
    outputs=gr.File(),
    title="Extractor de Nombres",
    description="Sube un archivo .docx y extrae los nombres de las personas usando NLP con SpaCy. Descarga el resultado en un archivo Excel."
)

# Iniciar la aplicación
iface.launch()