"""BM25 from scratch vs rank_bm25: correctness and speed on FastAPI docs chunks."""
import glob, math, re, statistics, sys, time, random, platform
from collections import Counter, defaultdict
import numpy, rank_bm25
from rank_bm25 import BM25Okapi

TOKEN = re.compile(r"\w+")
def tokenize(t): return TOKEN.findall(t.lower())

class BM25:
    """Okapi BM25 with an inverted index. idf='robertson' matches rank_bm25; 'lucene' is log(1+...)."""
    def __init__(self, docs, k1=1.5, b=0.75, idf="lucene", epsilon=0.25):
        self.k1, self.b, self.N = k1, b, len(docs)
        self.dl = [len(d) for d in docs]
        self.avgdl = sum(self.dl) / self.N
        self.post = defaultdict(list)           # term -> [(doc_id, tf)]
        for i, d in enumerate(docs):
            for t, f in Counter(d).items():
                self.post[t].append((i, f))
        self.idf = {}
        for t, p in self.post.items():
            n = len(p)
            self.idf[t] = (math.log(1 + (self.N - n + 0.5) / (n + 0.5)) if idf == "lucene"
                           else math.log((self.N - n + 0.5) / (n + 0.5)))
        if idf == "robertson":                   # rank_bm25 floors negative idf
            floor = epsilon * sum(self.idf.values()) / len(self.idf)
            self.idf = {t: (v if v >= 0 else floor) for t, v in self.idf.items()}
    def scores(self, query):
        s = defaultdict(float)
        for t in query:   # repeated query terms count repeatedly, like rank_bm25
            idf = self.idf.get(t)
            if idf is None: continue
            for i, f in self.post[t]:
                s[i] += idf * f * (self.k1 + 1) / (f + self.k1 * (1 - self.b + self.b * self.dl[i] / self.avgdl))
        return s
    def top(self, query, k=10):
        s = self.scores(query)
        return sorted(s, key=lambda i: (-s[i], i))[:k]

# --- corpus: split each file into sections at Markdown headings ---
chunks, heads = [], []
for fn in sorted(glob.glob("experiments/corpus/*.md")):
    cur, title = [], ""
    for line in open(fn, encoding="utf-8").read().splitlines() + ["# END"]:
        if re.match(r"#{1,4} ", line):
            body = " ".join(cur).strip()
            if len(body) > 200:
                chunks.append(title + " " + body); heads.append(title)
            cur, title = [], line.lstrip("# ").strip()
        else: cur.append(line)
docs = [tokenize(c) for c in chunks]
print(f"python {platform.python_version()} numpy {numpy.__version__} rank_bm25 {getattr(rank_bm25,'__version__','0.2.2')}")
print(f"chunks={len(docs)} avg_len={sum(map(len,docs))/len(docs):.1f} tokens, vocab={len(set(t for d in docs for t in d))}")

random.seed(0)
queries = [tokenize(h) for h in heads if len(tokenize(h)) >= 2]
queries = random.sample(queries, 300)
# longer natural-ish queries: 8 random tokens from a random chunk
queries_long = [random.sample(d, 8) for d in random.sample(docs, 200)]
print(f"queries: {len(queries)} heading (2+ tokens), {len(queries_long)} long (8 tokens)")

# --- 1. correctness: same idf formula as rank_bm25 ---
ref = BM25Okapi(docs, k1=1.5, b=0.75)
mine_r = BM25(docs, idf="robertson")
maxdiff, top10_same, top1_same = 0.0, 0, 0
ties_at_cut = diffs = 0
allq = queries + queries_long
for q in allq:
    rs = ref.get_scores(q); ms = mine_r.scores(q)
    for i in range(len(docs)):
        maxdiff = max(maxdiff, abs(rs[i] - ms.get(i, 0.0)))
    rtop = [i for i in numpy.argsort(-rs, kind="stable")[:10] if rs[i] > 0]  # drop zero-score filler
    mtop = mine_r.top(q, 10)
    top10_same += set(rtop) == set(mtop); top1_same += rtop[0] == mtop[0]
    if set(rtop) != set(mtop):   # is the difference only a tie at the rank-10 boundary?
        ties_at_cut += abs(rs[rtop[-1]] - ms.get(mtop[-1], 0.0)) < 1e-9; diffs += 1
print(f"\n[correctness, robertson idf, {len(allq)} queries x {len(docs)} docs]")
print(f"max abs score difference: {maxdiff:.3e}")
print(f"top-10 set identical: {top10_same}/{len(allq)}; top-1 identical: {top1_same}/{len(allq)}")
print(f"of {diffs} queries with a different top-10 set, {ties_at_cut} differ only by a score tie at rank 10")

# --- 2. lucene idf vs rank_bm25 ---
mine_l = BM25(docs, idf="lucene")
t10 = t1 = 0
for q in allq:
    rs = ref.get_scores(q)
    rtop = [i for i in numpy.argsort(-rs, kind="stable")[:10] if rs[i] > 0]; mtop = mine_l.top(q, 10)
    t10 += len(set(rtop) & set(mtop)) / max(len(rtop), 1); t1 += rtop[0] == mtop[0]
print(f"\n[lucene idf vs rank_bm25] mean top-10 overlap {t10/len(allq):.3f}; top-1 identical {t1}/{len(allq)}")

# --- 3. speed ---
def timeit(fn, qs, reps=3):
    best = None
    for _ in range(reps):
        t = time.perf_counter()
        for q in qs: fn(q)
        d = (time.perf_counter() - t) / len(qs) * 1000
        best = d if best is None else min(best, d)
    return best
def ref_top(q): return numpy.argsort(-ref.get_scores(q))[:10]
print("\n[speed: ms per query, best of 3 passes]")
print("| query set | rank_bm25 | from scratch (inverted index) | speedup |")
print("|---|---:|---:|---:|")
for name, qs in [("heading", queries), ("long", queries_long)]:
    a = timeit(ref_top, qs); c = timeit(lambda q: mine_r.top(q, 10), qs)
    print(f"| {name} | {a:.3f} | {c:.3f} | {a/c:.1f}x |")
t = time.perf_counter(); BM25Okapi(docs); a = time.perf_counter() - t
t = time.perf_counter(); BM25(docs); c = time.perf_counter() - t
print(f"\nindex build: rank_bm25 {a*1000:.0f} ms, from scratch {c*1000:.0f} ms")