"""
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_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
#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 streamlit as st

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 = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L12-v2")
    persist_directory = "./persist_directory"
    db = Chroma(persist_directory=persist_directory, embedding_function=embeddings.embed_documents)
    vectorstore = Chroma.from_documents(textos[:10], embeddings, persist_directory="./chroma_db")    
    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
    qa_chain=src.model_load.load_model()
    print("Lanzando interfaz")
    # Título de la aplicación
    st.title("Interfaz Sencilla")
    
    # Entrada de texto
    user_input = st.text_input("Escribe algo:")
    
    # Botón para obtener la respuesta
    if st.button("Obtener respuesta"):
        st.write(src.model_load.ask(user_input, retriever,qa_chain))