"""BM25 on the rag-context-lab fixture (pure Python, no model download).

Same fixture, same whitespace tokens and substring matching as compare.py; only the scoring changes
from "count of matched query words" to BM25 (k1=1.5, b=0.75). Run next to fixture.json.
"""
import json
import math
from pathlib import Path

root = Path(__file__).resolve().parent
fixture = json.loads((root / 'fixture.json').read_text(encoding='utf-8'))
docs = fixture['documents']


def tokens(text):
    return text.casefold().split()


def bm25_rank(question, include_title, k1=1.5, b=0.75):
    texts = {d['id']: ((d['title'] + '\n') if include_title else '') + d['body'] for d in docs}
    toks = {i: tokens(t) for i, t in texts.items()}
    avg = sum(len(t) for t in toks.values()) / len(toks)
    n = len(toks)
    scores = {}
    for i, words in toks.items():
        score = 0.0
        for term in set(tokens(question)):
            tf = sum(term in w for w in words)              # substring match, like compare.py
            if not tf:
                continue
            df = sum(any(term in w for w in ws) for ws in toks.values())
            idf = math.log(1 + (n - df + 0.5) / (df + 0.5))
            score += idf * tf * (k1 + 1) / (tf + k1 * (1 - b + b * len(words) / avg))
        if score:
            scores[i] = round(score, 3)
    return sorted(scores.items(), key=lambda kv: (-kv[1], kv[0]))


for mode in (False, True):
    print('title_and_body' if mode else 'body_only')
    for q in fixture['queries']:
        ranked = bm25_rank(q['text'], mode)
        top = ranked[0] if ranked else ('-', 0)
        ties = sum(1 for _, s in ranked if ranked and s == ranked[0][1])
        mark = 'hit' if top[0] == q['expected'] else ('no-answer' if q['expected'] is None else 'miss')
        print(f"  {q['id']} {q['text']:<10} expected={q['expected'] or '-':<3} top={top[0]:<3} score={top[1]:<6} ties={ties} {mark}")