# import faiss

# merged_index = faiss.read_index("/home/mshahidul/readctrl/data/vector_db/qwen_em/shard_0_en.faiss")
# for i in range(1, 2):
#     next_index = faiss.read_index(f"/home/mshahidul/readctrl/data/vector_db/qwen_em/shard_{i}_en.faiss")
#     merged_index.merge_from(next_index)

# faiss.write_index(merged_index, "/home/mshahidul/readctrl/data/vector_db/qwen_em/full_wikipedia_index.faiss")

import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "2"

import faiss
import numpy as np
import torch
from sentence_transformers import SentenceTransformer

# 1. Configuration
model_id = "Qwen/Qwen3-Embedding-4B"
index_path = "/home/mshahidul/readctrl/data/vector_db/qwen_em/full_wikipedia_index.faiss"

# 2. Load the Index
print("Loading Index...")
index = faiss.read_index(index_path)
print(f"Index loaded successfully.")
print(f"Total vectors in index: {index.ntotal}")
print(f"Vector dimension: {index.d}")

# 3. Load Model for Querying
print("Loading model for query embedding...")
model = SentenceTransformer(
    model_id, 
    trust_remote_code=True, 
    device="cuda",
    model_kwargs={"torch_dtype": torch.bfloat16}
)

# 4. Perform a Search
query = "What is the capital of France?"
# We must encode the query using the same model
query_vector = model.encode([query], convert_to_numpy=True).astype('float32')

k = 5  # Number of nearest neighbors to find
distances, indices = index.search(query_vector, k)

# 5. Review Results
print("\n--- Search Results ---")
print(f"Query: {query}")
for i in range(k):
    print(f"Result {i+1}: Index ID {indices[0][i]}, Distance: {distances[0][i]:.4f}")

if indices[0][0] == -1:
    print("\nError: The search returned -1. This usually means the index is empty or improperly trained.")
else:
    print("\nSuccess: The index returned valid neighbors!")