"""minimal-retrieval-pipeline.py -- a complete retrieval pipeline in the standard library.

What it does: loads a corpus of arXiv abstracts (datasets/minimal-retrieval-pipeline.csv,
pulled from the arXiv export API on 2026-09-05), splits each abstract into overlapping
chunks, builds TF-IDF vectors (log-scaled term frequency, smoothed IDF, L2-normalised),
runs cosine search, and evaluates the search with the paper titles as queries: a query is
"answered" when a chunk of its own abstract is retrieved. Reports Recall@1/5/10 and MRR
for whole-abstract indexing versus chunked indexing.

Inputs:  datasets/minimal-retrieval-pipeline.csv  (columns: arxiv_id, published, title, abstract, url)
Run:     python code/minimal-retrieval-pipeline.py [path/to/csv]
Requires: Python 3.11+ standard library only. Swap `embed()` for a neural embedding model
          (for example sentence-transformers) and nothing else changes.
"""
from __future__ import annotations

import csv
import math
import re
import sys
from collections import Counter, defaultdict

TOKEN = re.compile(r"[a-z0-9]+(?:-[a-z0-9]+)*")
STOP = set("a an the of and or to in for on with by from as is are was were be been that this these those it its we our".split())


def tokenize(text: str) -> list[str]:
    return [t for t in TOKEN.findall(text.lower()) if t not in STOP and len(t) > 1]


def chunk(text: str, size: int = 60, overlap: int = 20) -> list[str]:
    """Fixed-size word windows with overlap. size/overlap are in words."""
    words = text.split()
    if len(words) <= size:
        return [text]
    out, start = [], 0
    while start < len(words):
        out.append(" ".join(words[start:start + size]))
        if start + size >= len(words):
            break
        start += size - overlap
    return out


class TfidfIndex:
    def __init__(self, docs: list[tuple[str, str]]):
        """docs: list of (doc_id, text). Several chunks may share a doc_id."""
        self.ids = [d for d, _ in docs]
        tfs = [Counter(tokenize(t)) for _, t in docs]
        df: Counter = Counter()
        for tf in tfs:
            df.update(tf.keys())
        n = len(docs)
        self.idf = {t: math.log((1 + n) / (1 + c)) + 1.0 for t, c in df.items()}
        self.vectors = [self._vec(tf) for tf in tfs]
        self.postings: dict[str, list[int]] = defaultdict(list)
        for i, tf in enumerate(tfs):
            for t in tf:
                self.postings[t].append(i)

    def _vec(self, tf: Counter) -> dict[str, float]:
        v = {t: (1 + math.log(c)) * self.idf.get(t, 1.0) for t, c in tf.items()}
        norm = math.sqrt(sum(x * x for x in v.values())) or 1.0
        return {t: x / norm for t, x in v.items()}

    def embed(self, text: str) -> dict[str, float]:
        return self._vec(Counter(tokenize(text)))

    def search(self, query: str, k: int = 10) -> list[tuple[str, float]]:
        q = self.embed(query)
        candidates = set()
        for t in q:
            candidates.update(self.postings.get(t, ()))
        scored = []
        for i in candidates:
            v = self.vectors[i]
            scored.append((sum(w * v.get(t, 0.0) for t, w in q.items()), i))
        scored.sort(reverse=True)
        # collapse chunks: best chunk per document, keep document order
        seen, out = set(), []
        for s, i in scored:
            d = self.ids[i]
            if d in seen:
                continue
            seen.add(d)
            out.append((d, s))
            if len(out) == k:
                break
        return out


def evaluate(index: TfidfIndex, queries: list[tuple[str, str]], k: int = 10) -> dict[str, float]:
    hits = {1: 0, 5: 0, 10: 0}
    rr = 0.0
    for gold, q in queries:
        ranked = [d for d, _ in index.search(q, k)]
        if gold in ranked:
            r = ranked.index(gold) + 1
            rr += 1.0 / r
            for kk in hits:
                if r <= kk:
                    hits[kk] += 1
    n = len(queries)
    return {"recall@1": hits[1] / n, "recall@5": hits[5] / n, "recall@10": hits[10] / n, "mrr": rr / n}


def main(path: str) -> None:
    with open(path, encoding="utf-8") as f:
        rows = list(csv.DictReader(f))
    queries = [(r["arxiv_id"], r["title"]) for r in rows]

    whole = TfidfIndex([(r["arxiv_id"], r["abstract"]) for r in rows])
    chunks = []
    for r in rows:
        chunks.extend((r["arxiv_id"], c) for c in chunk(r["abstract"]))
    chunked = TfidfIndex(chunks)

    print(f"corpus: {len(rows)} abstracts, {len(chunks)} chunks (60 words, 20 overlap)")
    for name, idx in (("whole-abstract", whole), ("chunked", chunked)):
        m = evaluate(idx, queries)
        print(f"{name:15s} " + "  ".join(f"{k}={v:.3f}" for k, v in m.items()))

    # a query that is not a title, to show what the ranking looks like
    q = "how much does long context cost compared with retrieval"
    print(f"\nquery: {q!r}")
    for d, s in chunked.search(q, 5):
        title = next(r["title"] for r in rows if r["arxiv_id"] == d)
        print(f"  {s:.3f}  {d}  {title[:70]}")


if __name__ == "__main__":
    main(sys.argv[1] if len(sys.argv) > 1 else "datasets/minimal-retrieval-pipeline.csv")