import os
from sentence_transformers import SentenceTransformer
import pandas as pd
import chromadb
from chromadb.utils import embedding_functions
from scipy.spatial.distance import cosine
import gradio as gr
from tqdm import tqdm
import re

# Crear directorio para almacenamiento persistente
os.makedirs('./data_embeddings', exist_ok=True)

# Cargar el modelo de embeddings
model = SentenceTransformer("jinaai/jina-embeddings-v2-base-es", trust_remote_code=True)

# Ruta del archivo CSV
file_path = './v12_db_turismo.csv'

# Cargar la base de datos
df = pd.read_csv(file_path, encoding='latin-1')

# Función para limpiar y normalizar texto
def limpiar_texto(texto):
    texto = str(texto)
    texto = re.sub(r'[^a-zA-ZáéíóúñÁÉÍÓÚÑ0-9\s]', '', texto)  # Eliminar caracteres especiales
    texto = texto.lower()  # Convertir a minúsculas
    return texto.strip()

# Crear una columna combinada para los embeddings
df['text'] = df.apply(
    lambda x: limpiar_texto(
        f"Estado: {x['estado']}, "
        f"Municipio: {x['municipio']}, "
        f"Área natural o aventura: {x['area_natural_aventura']}, "
        f"Resumen: {x['resumen']}, "
        f"Descripción: {x['descripcion']}, "
        f"Ubicación: {x['ubicacion']}, "
        f"Link: {x['link']}, "
        f"Latitud: {x['latitud']}, Longitud: {x['longitud']}"
    ), axis=1
)

# Generar embeddings en batches de 64 con barra de progreso
print("Generando embeddings en lotes...")
df['embeddings'] = list(
    model.encode(
        df['text'].tolist(),  # Convertir la columna 'text' a una lista
        batch_size=64,        # Tamaño del batch
        show_progress_bar=True  # Mostrar barra de progreso
    )
)

# Crear identificadores únicos
df['ids'] = df.index.astype(str)

# Configurar cliente de ChromaDB persistente
chroma_client = chromadb.PersistentClient(path='./data_embeddings')
embedding_function = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="jinaai/jina-embeddings-v2-base-es")

db = chroma_client.get_or_create_collection(name='turismo_db', embedding_function=embedding_function)

# Agregar datos a la colección
db.add(
    ids=df['ids'].tolist(),
    embeddings=df['embeddings'].tolist(),
    documents=df['text'].tolist(),
    metadatas=df[['estado', 'municipio', 'area_natural_aventura', 'resumen',
                  'descripcion', 'ubicacion', 'link', 'latitud', 'longitud']].to_dict('records')
)

print(f"Registros en la colección: {db.count()}")

# Función para realizar consultas
def obtener_recomendaciones(preferencias, top_k=5):
    # Limpiar y generar embeddings de la consulta
    preferencias = limpiar_texto(preferencias)
    consulta_embedding = model.encode(preferencias, show_progress_bar=False)

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

    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]
        similitud_coseno = 1 - cosine(consulta_embedding, df['embeddings'].iloc[idx])
        metadata['distancia_euclidiana'] = dist
        metadata['similitud_coseno'] = similitud_coseno
        recomendaciones.append(metadata)

    # Ordenar por similitud coseno y distancia euclidiana
    recomendaciones = sorted(recomendaciones, key=lambda x: (x['distancia_euclidiana'], -x['similitud_coseno']))
    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."}])

    resultados = []
    for idx, sitio in enumerate(recomendaciones, 1):
        resultados.append({
            "#": idx,
            "Estado": sitio.get('estado', 'Desconocido'),
            "Municipio": sitio.get('municipio', 'No disponible'),
            "Área Natural o Aventura": sitio.get('area_natural_aventura', 'No disponible'),
            "Resumen": sitio.get('resumen', 'No disponible'),
            "Descripción": sitio.get('descripcion', 'No disponible'),
            "Ubicación": sitio.get('ubicacion', 'No disponible'),
            "Link": sitio.get('link', 'No disponible'),
            "Distancia Euclidiana": sitio.get('distancia_euclidiana', 'N/A'),
            "Similitud Coseno": sitio.get('similitud_coseno', 'N/A'),
        })

    return pd.DataFrame(resultados)

# Interfaz de Gradio
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., cañones, cascadas, áreas naturales"),
    outputs=gr.Dataframe(label="Resultados de Búsqueda"),
    title="Buscador de Turismo con Búsqueda Semántica Mejorada",
    description="Introduce tus preferencias para obtener resultados relevantes basados en similitud semántica y ordenados por relevancia."
)

if __name__ == "__main__":
    ui.launch(server_name="0.0.0.0", server_port=7860)