import os
from glob import glob
from sentence_transformers import SentenceTransformer
import gradio as gr
import numpy as np
import faiss

# Cargar el modelo de Sentence Transformers
model_name = "quora-distilbert-multilingual"
model = SentenceTransformer(model_name)

# Función para vectorizar un texto
def vectorize_text(text):
    return model.encode([text])[0]

# Obtener una lista de todos los archivos *.index en el directorio
archivos_indices = glob("*.index")

# Cargar todos los índices de Faiss
indices = []
nombres_archivos = []  # Lista para almacenar los nombres de los archivos

for archivo_index in archivos_indices:
    index = faiss.read_index(archivo_index)
    indices.append(index)
    nombres_archivos.append(os.path.splitext(archivo_index)[0])  # Obtener el nombre sin la extensión

# Función para predecir la similitud con cada índice
def predecir_similitud(texto):
    resultados = []
    umbral_similitud = 0.65

    # Vectorizar el texto (usando la misma función que antes)
    texto_vectorizado = vectorize_text(texto)

    for i, (index, nombre_archivo) in enumerate(zip(indices, nombres_archivos)):
        # Buscar similitud con Faiss
        texto_vectorizado_np = np.array([texto_vectorizado], dtype=np.float32)
        D, I = index.search(texto_vectorizado_np, 1)  # Buscar el vector más cercano
        similitud = 1 - D[0][0]  # Calcular similitud de coseno

        # Comparar con el umbral y proporcionar el resultado
        if similitud >= umbral_similitud:
            resultado = (nombre_archivo, similitud)
            resultados.append(resultado)

    # Ordenar la lista de resultados por similitud
    resultados_ordenados = sorted(resultados, key=lambda x: x[1], reverse=True)

    # Construir la cadena de resultados
    if resultados_ordenados:
        resultados_texto = '\n'.join([f"{nombre} (similitud:{sim:.2f})" for nombre, sim in resultados_ordenados])
    else:
        resultados_texto = "No se han encontrado resultados"

    return resultados_texto

# Crear la interfaz de Gradio con un botón de "Predecir"
iface = gr.Interface(
    fn=predecir_similitud,
    inputs=gr.inputs.Textbox(default=""),
    outputs=gr.outputs.Textbox(label="Argumentos"),  # Usar multiple=True
    live=False,
    title="Argumentos",
    description="Ingrese un texto y haga clic en 'Submit' para detectar Argumentos."
)

# Ejecutar la interfaz
iface.launch()