"""Minimal, auditable retrieval-rerank-citation baseline.

The default toy backend uses only Python's standard library. It is a protocol
smoke test, not a substitute for a trained embedding model or CrossEncoder.
Use --backend neural after installing requirements.txt for those components.
"""

from __future__ import annotations

import argparse
import hashlib
import math
import re
from collections import Counter
from dataclasses import dataclass
from typing import Iterable


SEED = 7
TOKEN_RE = re.compile(r"[a-z0-9]+")


@dataclass(frozen=True)
class Document:
    doc_id: str
    text: str


CORPUS = [
    Document("d1", "BM25 is a sparse lexical ranking function. The k1 parameter controls term-frequency saturation, while b controls document-length normalization."),
    Document("d2", "Dense retrieval independently encodes a query and each passage into fixed-size vectors, then ranks passages by dot product or cosine similarity."),
    Document("d3", "A cross-encoder jointly reads a query and one candidate passage. It is slower than a bi-encoder, so it commonly reranks only a small retrieved set."),
    Document("d4", "Citation evaluation should test whether cited evidence supports each claim and whether citations are irrelevant, not merely whether citation markers are present."),
    Document("d5", "Retrieval-augmented generation retrieves external passages and conditions generation on them, allowing knowledge to be traced to a non-parametric corpus."),
    Document("d6", "Activation checkpointing trades additional computation for lower training memory by recomputing selected forward activations during backward."),
]

QUERIES = [
    ("What controls BM25 term frequency saturation?", {"d1"}),
    ("Why use a cross encoder after retrieval?", {"d3"}),
    ("How should citations be evaluated?", {"d4"}),
    ("What is retrieval augmented generation?", {"d5"}),
]


def tokenize(text: str) -> list[str]:
    return TOKEN_RE.findall(text.lower())


class BM25:
    def __init__(self, documents: list[Document], k1: float = 1.2, b: float = 0.75):
        self.documents = documents
        self.k1 = k1
        self.b = b
        self.tokens = [tokenize(doc.text) for doc in documents]
        self.term_frequencies = [Counter(tokens) for tokens in self.tokens]
        self.avg_length = sum(map(len, self.tokens)) / len(self.tokens)
        self.document_frequency = Counter()
        for tokens in self.tokens:
            self.document_frequency.update(set(tokens))

    def score(self, query: str, index: int) -> float:
        n_docs = len(self.documents)
        doc_length = len(self.tokens[index])
        tf = self.term_frequencies[index]
        score = 0.0
        for term in set(tokenize(query)):
            frequency = tf[term]
            if frequency == 0:
                continue
            df = self.document_frequency[term]
            idf = math.log(1.0 + (n_docs - df + 0.5) / (df + 0.5))
            denominator = frequency + self.k1 * (
                1.0 - self.b + self.b * doc_length / self.avg_length
            )
            score += idf * frequency * (self.k1 + 1.0) / denominator
        return score

    def rank(self, query: str) -> list[tuple[str, float]]:
        results = [
            (doc.doc_id, self.score(query, index))
            for index, doc in enumerate(self.documents)
        ]
        return sorted(results, key=lambda item: (-item[1], item[0]))


def feature_hash_vector(text: str, dimensions: int = 256) -> list[float]:
    """Deterministic non-neural dense proxy made from word and character features."""
    tokens = tokenize(text)
    features = tokens + [
        f"c3:{token[index:index + 3]}"
        for token in tokens
        for index in range(max(0, len(token) - 2))
    ]
    vector = [0.0] * dimensions
    for feature in features:
        digest = hashlib.blake2b(
            f"{SEED}:{feature}".encode("utf-8"), digest_size=8
        ).digest()
        number = int.from_bytes(digest, "big")
        vector[number % dimensions] += 1.0 if (number >> 8) & 1 else -1.0
    norm = math.sqrt(sum(value * value for value in vector)) or 1.0
    return [value / norm for value in vector]


def dot(left: Iterable[float], right: Iterable[float]) -> float:
    return sum(a * b for a, b in zip(left, right))


def dense_proxy_rank(query: str, documents: list[Document]) -> list[tuple[str, float]]:
    query_vector = feature_hash_vector(query)
    results = [
        (doc.doc_id, dot(query_vector, feature_hash_vector(doc.text)))
        for doc in documents
    ]
    return sorted(results, key=lambda item: (-item[1], item[0]))


def reciprocal_rank_fusion(
    rankings: list[list[tuple[str, float]]], constant: int = 60
) -> list[tuple[str, float]]:
    fused: Counter[str] = Counter()
    for ranking in rankings:
        for rank, (doc_id, _) in enumerate(ranking, start=1):
            fused[doc_id] += 1.0 / (constant + rank)
    return sorted(fused.items(), key=lambda item: (-item[1], item[0]))


def lexical_rerank(
    query: str, candidate_ids: list[str], documents: dict[str, Document]
) -> list[tuple[str, float]]:
    """A transparent interaction-feature proxy, not a learned CrossEncoder."""
    query_tokens = tokenize(query)
    query_set = set(query_tokens)
    results = []
    for doc_id in candidate_ids:
        doc_tokens = tokenize(documents[doc_id].text)
        overlap = len(query_set & set(doc_tokens)) / max(1, len(query_set))
        ordered_pairs = set(zip(doc_tokens, doc_tokens[1:]))
        query_pairs = set(zip(query_tokens, query_tokens[1:]))
        proximity = len(query_pairs & ordered_pairs) / max(1, len(query_pairs))
        results.append((doc_id, overlap + 0.5 * proximity))
    return sorted(results, key=lambda item: (-item[1], item[0]))


class NeuralBackend:
    def __init__(self) -> None:
        from sentence_transformers import CrossEncoder, SentenceTransformer

        self.embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
        self.reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L6-v2")

    def dense_rank(
        self, query: str, documents: list[Document]
    ) -> list[tuple[str, float]]:
        embeddings = self.embedder.encode(
            [query] + [doc.text for doc in documents], normalize_embeddings=True
        )
        query_vector, document_vectors = embeddings[0], embeddings[1:]
        results = [
            (doc.doc_id, float(query_vector @ vector))
            for doc, vector in zip(documents, document_vectors)
        ]
        return sorted(results, key=lambda item: (-item[1], item[0]))

    def rerank(
        self, query: str, candidate_ids: list[str], documents: dict[str, Document]
    ) -> list[tuple[str, float]]:
        pairs = [(query, documents[doc_id].text) for doc_id in candidate_ids]
        scores = self.reranker.predict(pairs)
        results = [(doc_id, float(score)) for doc_id, score in zip(candidate_ids, scores)]
        return sorted(results, key=lambda item: (-item[1], item[0]))


def retrieval_metrics(
    rankings: list[list[str]], gold_sets: list[set[str]], k: int
) -> dict[str, float]:
    hits = 0
    reciprocal_ranks = []
    for ranking, gold in zip(rankings, gold_sets):
        hits += int(bool(set(ranking[:k]) & gold))
        first = next((i for i, doc_id in enumerate(ranking, 1) if doc_id in gold), None)
        reciprocal_ranks.append(0.0 if first is None else 1.0 / first)
    return {
        f"recall@{k}": hits / len(rankings),
        "mrr": sum(reciprocal_ranks) / len(reciprocal_ranks),
    }


def citation_metrics(
    predicted: dict[str, list[str]], gold_support: dict[str, set[str]], valid_ids: set[str]
) -> dict[str, float]:
    all_citations = [doc_id for citations in predicted.values() for doc_id in citations]
    valid = sum(doc_id in valid_ids for doc_id in all_citations)
    precise = sum(
        doc_id in gold_support.get(claim_id, set())
        for claim_id, citations in predicted.items()
        for doc_id in citations
    )
    supported_claims = sum(
        bool(set(predicted.get(claim_id, [])) & supports)
        for claim_id, supports in gold_support.items()
    )
    return {
        "citation_validity": valid / max(1, len(all_citations)),
        "citation_precision_oracle": precise / max(1, len(all_citations)),
        "citation_recall_oracle": supported_claims / max(1, len(gold_support)),
    }


def run_pipeline(backend: str, top_k: int) -> None:
    documents = {doc.doc_id: doc for doc in CORPUS}
    bm25 = BM25(CORPUS)
    neural = NeuralBackend() if backend == "neural" else None
    final_rankings: list[list[str]] = []
    predicted_citations: dict[str, list[str]] = {}
    gold_support: dict[str, set[str]] = {}

    for query_index, (query, gold) in enumerate(QUERIES, start=1):
        sparse = bm25.rank(query)
        dense = neural.dense_rank(query, CORPUS) if neural else dense_proxy_rank(query, CORPUS)
        fused = reciprocal_rank_fusion([sparse, dense])
        candidates = [doc_id for doc_id, _ in fused[:top_k]]
        reranked = (
            neural.rerank(query, candidates, documents)
            if neural
            else lexical_rerank(query, candidates, documents)
        )
        ranking = [doc_id for doc_id, _ in reranked]
        final_rankings.append(ranking)
        claim_id = f"q{query_index}"
        predicted_citations[claim_id] = ranking[:1]
        gold_support[claim_id] = gold
        print(f"{claim_id} query={query!r}")
        print(f"  bm25_top3={[doc_id for doc_id, _ in sparse[:3]]}")
        print(f"  dense_top3={[doc_id for doc_id, _ in dense[:3]]}")
        print(f"  reranked={ranking}")

    print("retrieval", retrieval_metrics(final_rankings, [gold for _, gold in QUERIES], 3))
    print("citations", citation_metrics(predicted_citations, gold_support, set(documents)))


def run_checks() -> None:
    bm25 = BM25(CORPUS)
    assert bm25.rank("BM25 term frequency saturation")[0][0] == "d1"
    vector = feature_hash_vector("deterministic embedding")
    assert len(vector) == 256
    assert math.isclose(dot(vector, vector), 1.0, abs_tol=1e-9)
    fused = reciprocal_rank_fusion([
        [("d1", 1.0), ("d2", 0.5)],
        [("d2", 1.0), ("d1", 0.5)],
    ])
    assert [doc_id for doc_id, _ in fused] == ["d1", "d2"]
    diagnostic = citation_metrics(
        {"c1": ["d1", "d6"], "c2": ["missing"]},
        {"c1": {"d1"}, "c2": {"d4"}},
        {doc.doc_id for doc in CORPUS},
    )
    assert math.isclose(diagnostic["citation_validity"], 2 / 3)
    assert math.isclose(diagnostic["citation_precision_oracle"], 1 / 3)
    assert math.isclose(diagnostic["citation_recall_oracle"], 1 / 2)
    print("check-only: passed")
    print("bm25_top1=d1 feature_vector_shape=[256]")
    print(f"diagnostic_citations={diagnostic}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--backend", choices=("toy", "neural"), default="toy")
    parser.add_argument("--top-k", type=int, default=4)
    parser.add_argument("--check-only", action="store_true")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    print(f"seed={SEED} corpus_size={len(CORPUS)} backend={args.backend}")
    run_checks()
    if not args.check_only:
        run_pipeline(args.backend, args.top_k)


if __name__ == "__main__":
    main()