import numpy as np
from sklearn.neighbors import NearestNeighbors

def compute_ground_truth(corpus_vectors, query_vectors, k=10):
    nn = NearestNeighbors(n_neighbors=k, algorithm="brute", metric="cosine")
    nn.fit(corpus_vectors)
    distances, indices = nn.kneighbors(query_vectors)
    return indices  # shape (n_queries, k)

def recall_at_k(retrieved_ids, ground_truth_ids):
    """Both: list of doc-id arrays, length n_queries; each row top-k retrieved."""
    recalls = []
    for retr, gt in zip(retrieved_ids, ground_truth_ids):
        recalls.append(len(set(retr) & set(gt)) / len(gt))
    return float(np.mean(recalls))

from qdrant_client import QdrantClient, models

client = QdrantClient(url="http://localhost:6333")
COLLECTION = "docs"  # points upserted with id == corpus index (to match ground truth)

def qdrant_search(query_vec, ef, k=10):
    # exact=False keeps the query on the HNSW path; exact=True is the oracle.
    resp = client.query_points(
        collection_name=COLLECTION,
        query=query_vec,
        search_params=models.SearchParams(hnsw_ef=ef, exact=False),
        limit=k,
        with_payload=False,
    )
    return [p.id for p in resp.points]

ef_values = [16, 32, 64, 128, 256]
results = []
for ef in ef_values:
    retrieved = [qdrant_search(q, ef=ef, k=10) for q in query_vectors]
    recall = recall_at_k(retrieved, ground_truth)
    results.append({"ef": ef, "recall": recall})
    # Pair with measure_latency() from Step 4 to capture p95 at this ef.

# Find smallest ef that hits the recall target.
for r in results:
    if r["recall"] >= 0.95:
        print(f"smallest ef hitting target: ef={r['ef']} -> recall {r['recall']:.3f}")
        break

import time

def measure_latency(queries, engine_search_fn):
    times = []
    for q in queries:
        t0 = time.perf_counter()
        engine_search_fn(q, k=10)
        times.append((time.perf_counter() - t0) * 1000)
    times.sort()
    return {
        "p50": times[len(times) // 2],
        "p95": times[int(len(times) * 0.95)],
        "p99": times[int(len(times) * 0.99)],
    }

def test_latency_under_budget():
    lat = measure_latency(queries, qdrant_search)
    assert lat["p95"] < 50, f"p95 {lat['p95']:.1f}ms exceeds 50ms budget"

def test_recall_holds_after_embedding_upgrade():
    # Pre-upgrade baseline
    old_recall = baseline_recall_for("text-embedding-3-small")

    # Re-embed corpus + queries
    new_corpus = embed_with("text-embedding-3-large")
    new_queries = embed_with("text-embedding-3-large")
    new_gt = compute_ground_truth(new_corpus, new_queries)

    # Re-build vector index
    rebuild_index(new_corpus)

    # Measure
    new_recall = recall_at_k(retrieve_topk(new_queries, k=10), new_gt)
    delta = new_recall - old_recall
    # Allow up to 5% recall drop on engineering basis; flag if larger
    assert delta >= -0.05, f"Recall dropped {delta:.3f} after embedding upgrade"

def test_filtered_search_recall():
    # Filter: only docs with category=shoes
    filter_fn = lambda doc: doc["category"] == "shoes"
    filtered_corpus = [v for v, d in zip(corpus_vectors, corpus_metadata) if filter_fn(d)]
    filtered_gt = compute_ground_truth(filtered_corpus, query_vectors)

    retrieved = engine_search(query_vectors, filter={"category": "shoes"}, k=10)
    recall = recall_at_k(retrieved, filtered_gt)
    assert recall >= 0.90

def test_recall_after_incremental_inserts():
    # Initial index
    initial_recall = recall_at_k(engine_search(queries, k=10), ground_truth)

    # Insert 10% more docs
    for v, m in new_docs:
        engine.insert(v, m)

    # Re-measure (without rebuild)
    new_recall = recall_at_k(engine_search(queries, k=10), ground_truth_extended)
    assert new_recall >= initial_recall - 0.05