import argparse
from statistics import median
from time import perf_counter

import faiss
import numpy as np

DIMENSIONS = 96
TOP_K = 10


def unit_vectors(generator, count):
    vectors = generator.standard_normal((count, DIMENSIONS), dtype=np.float32)
    faiss.normalize_L2(vectors)
    return vectors


def build(index, vectors, label):
    start = perf_counter()
    index.add(vectors)
    elapsed = perf_counter() - start
    size = faiss.serialize_index(index).nbytes / 1024**2
    print(f"{label}: {elapsed:.2f}s to build, {size:.1f} MiB serialized")


def measure(index, questions):
    index.search(questions[:5], TOP_K)
    results, times = [], []
    for question in questions:
        start = perf_counter()
        _, ids = index.search(question.reshape(1, -1), TOP_K)
        times.append((perf_counter() - start) * 1000)
        results.append(ids[0])
    return results, median(times)


def compare(count, queries, seed):
    faiss.omp_set_num_threads(1)
    generator = np.random.default_rng(seed)
    vectors = unit_vectors(generator, count)
    questions = unit_vectors(generator, queries)

    exact = faiss.IndexFlatIP(DIMENSIONS)
    graph = faiss.IndexHNSWFlat(DIMENSIONS, 16, faiss.METRIC_INNER_PRODUCT)
    graph.hnsw.efConstruction = 80
    build(exact, vectors, "Exact")
    build(graph, vectors, "HNSW")

    baseline, exact_time = measure(exact, questions)
    print(f"Exact search: {exact_time:.3f} ms median per query")
    for width in (10, 32, 96):
        graph.hnsw.efSearch = width
        results, elapsed = measure(graph, questions)
        hits = sum(len(set(expected) & set(observed)) for expected, observed in zip(baseline, results))
        print(f"HNSW efSearch={width:2}: recall@{TOP_K} {hits / (queries * TOP_K):.1%}, "
              f"{elapsed:.3f} ms median per query")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Compare exact and graph-based vector search on the same inputs.")
    parser.add_argument("--vectors", type=int, default=30_000)
    parser.add_argument("--queries", type=int, default=100)
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()
    if args.vectors < TOP_K or args.queries < 1 or args.seed < 0:
        parser.error(f"Use at least {TOP_K} vectors, one query, and a nonnegative seed.")
    compare(args.vectors, args.queries, args.seed)