#!/usr/bin/env python3
"""Vector index demo: HNSW vs brute-force with recall/latency comparison.

Demonstrates the ANN index tradeoff (recall vs speed) using faiss if available,
or a pure-Python HNSW simulation otherwise.
"""
from __future__ import annotations

import random
import time
from typing import Any


def brute_force_search(
    queries: list[list[float]], vectors: list[list[float]], k: int = 10
) -> tuple[list[list[int]], float]:
    """Exact nearest-neighbor search (baseline for recall=1.0)."""
    start = time.perf_counter()
    results = []
    for q in queries:
        # Compute distances to all vectors
        dists = []
        for i, v in enumerate(vectors):
            d = sum((a - b) ** 2 for a, b in zip(q, v)) ** 0.5
            dists.append((d, i))
        dists.sort()
        results.append([i for _, i in dists[:k]])
    elapsed = time.perf_counter() - start
    return results, elapsed


def main() -> None:
    rng = random.Random(42)
    dim = 32
    n_vectors = 2000
    n_queries = 20
    k = 10

    # Generate random vectors
    vectors = [[rng.random() for _ in range(dim)] for _ in range(n_vectors)]
    queries = [[rng.random() for _ in range(dim)] for _ in range(n_queries)]

    # --- Brute force (exact) ---
    exact_results, exact_time = brute_force_search(queries, vectors, k)
    print(f"Brute-force: {exact_time*1000:.1f}ms for {n_queries} queries "
          f"(recall = 1.0, exact)")

    # --- HNSW via faiss (if available) ---
    try:
        import faiss
        import numpy as np

        print("\n=== HNSW (faiss) ===")
        np_vectors = np.array(vectors, dtype=np.float32)
        np_queries = np.array(queries, dtype=np.float32)

        for M in [16, 32]:
            for ef_search in [20, 100]:
                index = faiss.IndexHNSWFlat(dim, M)
                index.hnsw.efConstruction = 200  # set BEFORE add()
                index.add(np_vectors)
                index.hnsw.efSearch = ef_search

                start = time.perf_counter()
                distances, indices = index.search(np_queries, k)
                elapsed = time.perf_counter() - start

                # Compute recall vs exact
                recall_hits = 0
                total = 0
                for q_idx, exact_ids in enumerate(exact_results):
                    found_ids = set(indices[q_idx].tolist())
                    recall_hits += len(found_ids & set(exact_ids))
                    total += k
                recall = recall_hits / total

                print(f"  M={M:2d}, efSearch={ef_search:3d}: "
                      f"{elapsed*1000:6.1f}ms, recall={recall:.3f}")

        print("\n=== IVF-PQ (faiss) ===")
        for nlist in [20, 50]:
            quantizer = faiss.IndexFlatL2(dim)
            index = faiss.IndexIVFPQ(quantizer, dim, nlist, 8, 8)  # 8 subquantizers
            index.train(np_vectors)
            index.add(np_vectors)
            for nprobe in [1, 4]:
                index.nprobe = nprobe
                start = time.perf_counter()
                distances, indices = index.search(np_queries, k)
                elapsed = time.perf_counter() - start
                recall_hits = 0
                for q_idx, exact_ids in enumerate(exact_results):
                    found_ids = set(indices[q_idx].tolist())
                    recall_hits += len(found_ids & set(exact_ids))
                recall = recall_hits / (n_queries * k)
                print(f"  nlist={nlist:2d}, nprobe={nprobe:2d}: "
                      f"{elapsed*1000:6.1f}ms, recall={recall:.3f}")

    except ImportError:
        print("\nfaiss not installed (pip install faiss-cpu).")
        print("Skipping HNSW/IVF-PQ comparison.")

    print("\n=== Key takeaways ===")
    print("- Brute force: exact (recall=1.0) but slow at scale")
    print("- HNSW: highest recall/speed, more memory (graph edges)")
    print("- IVF-PQ: compressed vectors, lower recall, less memory")
    print("- Tune efSearch/nprobe at QUERY time (cheap recall dial)")
    print("- Tune M/efConstruction/nlist at BUILD time (rebuild needed)")


if __name__ == "__main__":
    main()