#!/usr/bin/env python3
"""Vector DB retrieval demo: hybrid search (dense + BM25) with RRF fusion.

Demonstrates the two retrieval signals and how reciprocal rank fusion combines
them for better recall than pure vector search.
"""
from __future__ import annotations

import math
import re
from collections import Counter


# --- Dense (vector) retrieval via cosine similarity ---

def cosine_sim(a: list[float], b: list[float]) -> float:
    dot = sum(x * y for x, y in zip(a, b))
    na = math.sqrt(sum(x * x for x in a))
    nb = math.sqrt(sum(y * y for y in b))
    return dot / (na * nb) if na and nb else 0.0


# --- BM25 (sparse/keyword) retrieval ---

def tokenize(text: str) -> list[str]:
    return re.findall(r'\w+', text.lower())


def bm25_score(query_terms: list[str], doc_terms: list[str], k1: float = 1.5, b: float = 0.75) -> float:
    """Simplified BM25 for a single document (no corpus IDF for brevity)."""
    doc_freq = Counter(doc_terms)
    score = 0.0
    for term in query_terms:
        tf = doc_freq.get(term, 0)
        if tf == 0:
            continue
        # Simplified: assume term is rare (high IDF)
        idf = 1.0
        score += idf * (tf * (k1 + 1)) / (tf + k1)
    return score


# --- Reciprocal Rank Fusion ---

def reciprocal_rank_fusion(
    dense_ranking: list[int],
    sparse_ranking: list[int],
    k: int = 60,
) -> list[tuple[int, float]]:
    """Fuse two ranked lists using RRF: score = sum(1/(k + rank))."""
    scores: dict[int, float] = {}
    for rank, doc_id in enumerate(dense_ranking):
        scores[doc_id] = scores.get(doc_id, 0.0) + 1.0 / (k + rank + 1)
    for rank, doc_id in enumerate(sparse_ranking):
        scores[doc_id] = scores.get(doc_id, 0.0) + 1.0 / (k + rank + 1)
    return sorted(scores.items(), key=lambda x: -x[1])


def main() -> None:
    # Small corpus (5 docs)
    documents = [
        "vector databases store embeddings for similarity search",
        "PostgreSQL with pgvector extension supports vector search",
        "HNSW is a graph-based approximate nearest neighbor index",
        "BM25 is a sparse keyword retrieval algorithm",
        "hybrid search combines dense and sparse retrieval signals",
    ]
    query = "hybrid vector keyword search"

    # Simulate dense ranking with random-but-plausible vectors
    # (In production: embed with a real model)
    dense_scores = [cosine_sim(
        [hash(q) % 10 + 1 for q in tokenize(query)][:8] or [1],
        [hash(t) % 10 + 1 for t in tokenize(doc)][:8] or [1],
    ) for doc in documents]

    dense_ranking = sorted(range(len(documents)), key=lambda i: -dense_scores[i])

    # Sparse ranking with BM25
    q_terms = tokenize(query)
    sparse_scores = [bm25_score(q_terms, tokenize(doc)) for doc in documents]
    sparse_ranking = sorted(range(len(documents)), key=lambda i: -sparse_scores[i])

    print("=== Dense-only ranking ===")
    for rank, i in enumerate(dense_ranking):
        print(f"  {rank+1}. (score={dense_scores[i]:.3f}) {documents[i][:60]}")

    print("\n=== Sparse-only ranking (BM25) ===")
    for rank, i in enumerate(sparse_ranking):
        print(f"  {rank+1}. (score={sparse_scores[i]:.3f}) {documents[i][:60]}")

    fused = reciprocal_rank_fusion(dense_ranking, sparse_ranking)

    print("\n=== Hybrid (RRF fused) ranking ===")
    for rank, (doc_id, score) in enumerate(fused):
        print(f"  {rank+1}. (RRF={score:.4f}) {documents[doc_id][:60]}")

    print("\n=== Key insight ===")
    print("Hybrid search recovers exact-token matches (BM25) that dense")
    print("embeddings miss, and paraphrases (dense) that BM25 misses.")


if __name__ == "__main__":
    main()