#!/usr/bin/env python3
"""Hybrid Vector + Keyword Search with Reciprocal Rank Fusion (RRF).

Demonstrates: BM25 retrieval, dense vector retrieval, RRF fusion, and
comparison of single-method vs hybrid results.
"""

from __future__ import annotations

import math
from dataclasses import dataclass


@dataclass
class Document:
    id: str
    text: str
    score: float = 0.0


# ── BM25 (keyword) retriever ────────────────────────────────────────────

class BM25Retriever:
    """Simple BM25 implementation (Robertson et al.)."""

    def __init__(self, corpus: list[str], k1: float = 1.5, b: float = 0.75):
        self.corpus = corpus
        self.k1 = k1
        self.b = b
        self.doc_lens = [len(doc.split()) for doc in corpus]
        self.avg_dl = sum(self.doc_lens) / max(len(corpus), 1)
        self.N = len(corpus)
        # Precompute IDF
        self.idf: dict[str, float] = {}
        for doc in corpus:
            for term in set(doc.lower().split()):
                self.idf[term] = self.idf.get(term, 0) + 1
        for term, df in self.idf.items():
            self.idf[term] = math.log((self.N - df + 0.5) / (df + 0.5) + 1)

    def search(self, query: str, top_k: int = 20) -> list[Document]:
        query_terms = query.lower().split()
        results = []
        for i, doc in enumerate(self.corpus):
            score = 0.0
            doc_terms = doc.lower().split()
            tf = {t: doc_terms.count(t) / max(len(doc_terms), 1) for t in set(doc_terms)}
            for term in query_terms:
                if term in self.idf:
                    t = tf.get(term, 0)
                    score += self.idf[term] * (
                        t * (self.k1 + 1) / (t + self.k1 * (1 - self.b + self.b * self.doc_lens[i] / self.avg_dl))
                    )
            if score > 0:
                results.append(Document(id=str(i), text=doc, score=score))
        results.sort(key=lambda d: d.score, reverse=True)
        return results[:top_k]


# ── Dense retriever (mock) ──────────────────────────────────────────────

def dense_search(query: str, corpus: list[str], top_k: int = 20) -> list[Document]:
    """Mock dense retriever: word-overlap + substring scoring as a stand-in."""
    results = []
    for i, doc in enumerate(corpus):
        ql = query.lower()
        dl = doc.lower()
        # Substring bonus for conceptual matches
        score = sum(1 for w in ql.split() if w in dl) * 0.5
        score += sum(1 for w in dl.split() if w in ql) * 0.3
        if ql in dl:
            score += 2.0
        if score > 0:
            results.append(Document(id=str(i), text=doc, score=score))
    results.sort(key=lambda d: d.score, reverse=True)
    return results[:top_k]


# ── Reciprocal Rank Fusion ──────────────────────────────────────────────

def reciprocal_rank_fusion(ranked_lists: list[list[Document]], k: int = 60) -> list[Document]:
    """Fuse multiple ranked lists into one using RRF."""
    fusion_scores: dict[str, tuple[float, Document]] = {}
    for ranked in ranked_lists:
        for rank, doc in enumerate(ranked):
            rrf = 1.0 / (k + rank + 1)
            if doc.id in fusion_scores:
                prev_score, prev_doc = fusion_scores[doc.id]
                fusion_scores[doc.id] = (prev_score + rrf, prev_doc)
            else:
                fusion_scores[doc.id] = (rrf, doc)
    fused = sorted(fusion_scores.values(), key=lambda x: x[0], reverse=True)
    return [doc for _, doc in fused]


# ── Demo ────────────────────────────────────────────────────────────────

CORPUS = [
    "Python is a popular programming language for data science and machine learning.",
    "JavaScript runs in web browsers and is used for frontend development.",
    "PostgreSQL is a powerful open-source relational database system.",
    "Docker containers package applications with their dependencies.",
    "Machine learning models require careful evaluation and testing.",
    "Web development with Python often uses Django or FastAPI frameworks.",
    "SQL databases like PostgreSQL support complex queries and transactions.",
    "Frontend frameworks include React, Vue, and Angular for building UIs.",
    "Data science pipelines often involve Python, pandas, and Jupyter notebooks.",
    "Container orchestration with Kubernetes manages Docker deployments at scale.",
]

if __name__ == "__main__":
    bm25 = BM25Retriever(CORPUS)
    queries = [
        ("web development Python", "exact phrase + concept"),
        ("SQL database", "exact term match"),
        ("machine learning evaluation", "conceptual query"),
    ]

    for query, kind in queries:
        print(f"\n{'='*60}")
        print(f"Query: '{query}' ({kind})")

        bm25_results = bm25.search(query, top_k=5)
        dense_results = dense_search(query, CORPUS, top_k=5)

        print("\n  BM25 only:")
        for r, doc in enumerate(bm25_results, 1):
            print(f"    {r}. {doc.text[:70]}...")

        print("\n  Dense only:")
        for r, doc in enumerate(dense_results, 1):
            print(f"    {r}. {doc.text[:70]}...")

        hybrid = reciprocal_rank_fusion([bm25_results, dense_results])
        print("\n  Hybrid (RRF):")
        for r, doc in enumerate(hybrid[:5], 1):
            print(f"    {r}. {doc.text[:70]}...")

    print("\n✓ Hybrid search demo complete.")