# Instalar bibliotecas necesarias (solo se ejecutará localmente en un Space)
import os
os.system("pip install gradio sentence-transformers pandas numpy scipy chromadb")

# Importar librerías necesarias
from sentence_transformers import SentenceTransformer
import pandas as pd
import chromadb
from chromadb.utils import embedding_functions
from datetime import datetime
import gradio as gr

# Configuración de variables
EMBEDDING_LLM = "jinaai/jina-embeddings-v2-base-es"  # Modelo de embedding utilizado

# Cargar el modelo de embeddings
model = SentenceTransformer(
    EMBEDDING_LLM,
    trust_remote_code=True,
)

# Ruta del archivo CSV (debe estar disponible en el directorio del Space)
file_path = './Diagrama_de_decisión completo.tsv'

# Cargar la base de datos, especificando la codificación
df = pd.read_csv(
    file_path,
    encoding='utf-8',
    sep='|'
)
df['ID'].astype(str).tolist()

# Crear una columna combinada para embeddings
df['text'] = df.apply(
    lambda x: f"Trámite: {x['Trámite']}; Descripción: {x['Descripción']}; Palabras clave: {x['Tag (palabras clave)']}",
    axis=1
)

# Generar embeddings
df['embeddings'] = model.encode(df['text'], batch_size=64, show_progress_bar=True).tolist()

# Configurar cliente de ChromaDB persistente (usando una ruta compatible con Spaces)
client_persistent = chromadb.PersistentClient(path='./data_embeddings')
embedding_function = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=EMBEDDING_LLM)

# Crear o cargar colección en ChromaDB
db = client_persistent.get_or_create_collection(
    name='diagrama_decision_db',
    embedding_function=embedding_function,
    metadata={
        'description': 'Trámites y servicios - Centro de Atención para el Bienestar',
        'created': str(datetime.now()),
        'hnsw:construction_ef': 120,
        'hnsw:search_ef': 120,
        'hnsw:M': 5
    }
)

# Agregar datos a la colección
db.add(
    ids=df['ID'].astype(str).tolist(),
    embeddings=df['embeddings'].tolist(),
    documents=df['text'].tolist(),
    metadatas=df[['Trámite', 'Descripción', 'Tag (palabras clave)']].to_dict('records')
)

print(f"Número de registros en la colección: {db.count()}")

# Función para realizar consultas
def obtener_recomendaciones(preferencias, top_k=5):
    # Convertir preferencias en embeddings
    consulta_embedding = model.encode(preferencias, show_progress_bar=False)

    # Realizar consulta en ChromaDB
    results = db.query(query_embeddings=[consulta_embedding], n_results=top_k)

    # Verificar si hay resultados
    if not results or not results["metadatas"] or not results["metadatas"][0]:
        return []

    recomendaciones = []
    for idx, metadata in enumerate(results["metadatas"][0]):
        dist = results["distances"][0][idx]
        metadata['distancia_euclidiana'] = dist
        recomendaciones.append(metadata)

    return recomendaciones

# Función para procesar resultados y generar una tabla
def procesar_recomendaciones(preferencias):
    recomendaciones = obtener_recomendaciones(preferencias)

    if not recomendaciones:
        return pd.DataFrame([{"Mensaje": "No se encontraron resultados para las preferencias proporcionadas."}])

    # Crear una tabla con las recomendaciones
    resultados = []
    for idx, tramite in enumerate(recomendaciones, 1):
        resultados.append({
            "#": idx,
            "Trámite": tramite.get('Trámite', 'Desconocido'),
            "Descripción": tramite.get('Descripción', 'No disponible'),
            "Palabras clave": tramite.get('Tag (palabras clave)', 'No disponible'),
            "Distancia Euclidiana": tramite.get('distancia_euclidiana', 'N/A')
        })

    # Convertir los resultados a un DataFrame
    return pd.DataFrame(resultados)

# Interfaz de Gradio con formato de tabla
def interfaz(preferencias):
    tabla_resultados = procesar_recomendaciones(preferencias)
    return tabla_resultados

ui = gr.Interface(
    fn=interfaz,
    inputs=gr.Textbox(label="Preferencias", placeholder="Escribe tus preferencias, e.g., consulta de trámites, palabras clave"),
    outputs=gr.Dataframe(label="Resultados de Búsqueda"),
    title="Buscador de Trámites y Servicios con Búsqueda Semántica",
    description="Introduce tus preferencias para obtener resultados relevantes basados en similitud semántica."
)

# Ejecutar la aplicación
ui.launch(server_name="0.0.0.0", server_port=7860)