"""
Este script realiza las siguientes tareas:
1. Carga archivos desde un directorio especificado utilizando la clase `Loader` del módulo `preprocess`.
2. Procesa los archivos mediante limpieza de texto y división en fragmentos.
3. Genera representaciones vectoriales de los textos utilizando `sentence-transformers`.
4. Almacena los vectores en una base de datos Chroma para su posterior recuperación.
5. Inicializa un modelo y ejecuta una interfaz para interactuar con los datos procesados.

Módulos utilizados:
- `preprocess`: Contiene la clase `Loader` para la carga y preprocesamiento de documentos.
- `vdb`: Se asume que gestiona la base de datos vectorial.
- `model_load`: Módulo para cargar el modelo de machine learning.
- `st`: Se asume que proporciona la interfaz de usuario.

Estructura del código:
1. Define el directorio de los archivos a procesar.
2. Carga los archivos y los procesa si el número de archivos es menor a 2.
3. Si hay múltiples archivos, los procesa en un bucle y concatena los fragmentos.
4. Genera embeddings utilizando `sentence-transformers/all-MiniLM-L12-v2`.
5. Almacena los embeddings en ChromaDB y configura un recuperador basado en similitud.
6. Carga el modelo de machine learning.
7. Inicia la interfaz de usuario..

"""
#from langchain.vectorstores import Chroma
from langchain_chroma import Chroma
from tqdm.auto import tqdm
#from chromadb.utils import embedding_functions
from src.preprocess import Loader
from src.vdb import EmbeddingGen
import src.model_load
import src.st
import os

if __name__=="__main__":
    # Definición de directorio
    archivo = [r"data/Calculo_Trascendentes_Tempranas_Zill_4t.pdf"]
    
    # Carga de archivos y procesamiento de texto
    if len(archivo) < 2:
        Load = Loader(archivo[0])
        documentos = Load.load_docs()
        textos_limpios = [Loader.limpiar_texto(texto=doc) for doc in documentos]
        textos = Loader.splitter(texto=textos_limpios, chunk_size=500, chunk_overlap=50)
    else:
        textos = []
        for i in range(len(archivo)):
            Load = Loader(archivo)
            documentos = Load.load_docs()
            textos_limpios = [Load.limpiar_texto(texto=doc) for doc in documentos]
            chunks = Load.splitter(texto=textos_limpios, chunk_size=500, chunk_overlap=50)
            textos.extend(chunks)
    
    # Generación de embeddings y almacenamiento en base de datos ChromaDB
    embeddings = EmbeddingGen("sentence-transformers/all-MiniLM-L12-v2")
    # Define la ruta donde deseas guardar el vector store
    persist_directory = './vectorstore_data'
    
    # Verifica si la carpeta de persistencia existe, si no la crea
    if not os.path.exists(persist_directory):
        os.makedirs(persist_directory)
    
        # Función para procesar documentos por lotes
    def process_in_batches(documents, batch_size):
        """Genera embeddings en lotes para evitar problemas de memoria."""
        for i in range(0, len(documents), batch_size):
            batch = [doc.page_content for doc in documents[i:i + batch_size]]
            embeddings_batch = embeddings.embed_documents(batch)  # Cambié encode por embed_documents
            yield embeddings_batch
    vectorstore = Chroma(persist_directory=persist_directory)
    # Uso de la función con tqdm
    for embeddings_batch in tqdm(process_in_batches(textos, batch_size=16), desc="Procesando lotes de documentos"):
        # Procesar embeddings_batch aquí
        # Aquí se asume que tienes algún método para agregar estos embeddings al vectorstore
        vectorstore.add_documents(embeddings_batch)

    
    print("Vectorizado terminado")
    
    # Convierte el vector store en un recuperador (retriever) de similaridad
    retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 3})
    
    # Guarda el vector store para que persista entre ejecuciones
    vectorstore.persist()
    print("Vectorizado terminado")
    retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 3})
    print("Carga del modelo")
    # Carga del modelo y ejecución de la interfaz
    src.model_load.load_model()
    print("Lanzando interfaz")
    src.st.interfaz()