"""Chunk size sweep on the real FastAPI English docs (run from experiments/: python chunk-size-sweep-rag/bench_size.py).
Retrieval = BM25 (k1=1.5, b=0.75). Queries = h2/h3 headings with >=300 chars of body and 2+ words (773).
Strict hit  : retrieved chunk lies >=50% inside the target section (same rule as the chunking post).
Lenient hit : retrieved chunk shares >=100 characters with the target section (does not penalise big chunks).
Strategies swept over SIZE in 200..2000: fixed windows (100 overlap, or 10% of size) and heading-aware packed chunks (+ heading path)."""
import re, glob, math, collections, statistics as st
TOK = re.compile(r"[a-z0-9_]+")
def toks(s): return TOK.findall(s.lower())
docs = {p: open(p, encoding='utf-8').read() for p in sorted(glob.glob('corpus/*.md'))}
H = re.compile(r"^(#{1,3}) (.+?)\s*(\{.*\})?$", re.M)
def sections(text):
    ms = list(H.finditer(text)); out = []; stack = []
    for i, m in enumerate(ms):
        lvl = len(m.group(1)); title = m.group(2).strip()
        stack = [s for s in stack if s[0] < lvl] + [(lvl, title)]
        end = ms[i + 1].start() if i + 1 < len(ms) else len(text)
        out.append((m.start(), end, lvl, title, " > ".join(t for _, t in stack)))
    return out
def fixed(text, size):
    over = size // 10; out, i = [], 0
    while i < len(text):
        out.append((i, min(len(text), i + size), '')); i += size - over
        if out[-1][1] == len(text): break
    return out
def paragraphs(text, size):
    out, start, cur = [], 0, 0
    for m in re.finditer(r"\n\s*\n", text):
        if m.end() - start > size and cur > start: out.append((start, cur, '')); start = cur
        cur = m.end()
    out.append((start, len(text), '')); return [c for c in out if text[c[0]:c[1]].strip()]
def heading_aware(text, size):
    out = []
    secs = sections(text) or [(0, len(text), 0, '', '')]
    if secs[0][0] > 0: secs = [(0, secs[0][0], 0, '', '')] + secs
    for s, e, _, _, path in secs:
        for a, b, _ in paragraphs(text[s:e], size): out.append((s + a, s + b, path))
    return out
def bm25_index(chunks):
    tf = [collections.Counter(c) for c in chunks]; N = len(chunks)
    df = collections.Counter(t for c in tf for t in c); avg = sum(map(len, chunks)) / N
    idf = {t: math.log(1 + (N - n + .5) / (n + .5)) for t, n in df.items()}
    inv = collections.defaultdict(list)
    for i, c in enumerate(tf):
        for t, f in c.items(): inv[t].append((i, f))
    L = [len(c) for c in chunks]
    def search(q, k=5):
        sc = collections.defaultdict(float)
        for t in set(toks(q)):
            for i, f in inv.get(t, ()):
                sc[i] += idf[t] * f * 2.5 / (f + 1.5 * (.25 + .75 * L[i] / avg))
        return sorted(sc, key=sc.get, reverse=True)[:k]
    return search
queries = [(p, s, e, title) for p, t in docs.items() for s, e, lvl, title, _ in sections(t)
           if lvl >= 2 and e - s >= 300 and len(toks(title)) >= 2]
print(f"docs={len(docs)} chars={sum(map(len, docs.values())):,} queries={len(queries)}")
print(f"{'strategy':28} {'size':>5} {'chunks':>7} {'avg len':>8} {'cross%':>7} {'strict@1':>9} {'strict@5':>9} {'MRR':>6} {'lenient@1':>10} {'lenient@5':>10}")
SIZES = [200, 400, 600, 800, 1200, 1600, 2000]
strategies = {'fixed (10% overlap)': lambda t, n: fixed(t, n),
              'heading-aware + path': lambda t, n: heading_aware(t, n)}
for name, fn in strategies.items():
    for size in SIZES:
        allc = []
        for p, t in docs.items():
            bounds = [s for s, *_ in sections(t)]
            for a, b, pre in fn(t, size):
                allc.append((p, a, b, (pre + '\n' + t[a:b]) if pre else t[a:b], any(a < x < b - 1 for x in bounds)))
        search = bm25_index([toks(c[3]) for c in allc])
        s1 = s5 = l1 = l5 = rr = 0
        for p, s, e, title in queries:
            res = search(title, 5)
            def ov(i): return min(e, allc[i][2]) - max(s, allc[i][1])
            st_r = [r for r, i in enumerate(res) if allc[i][0] == p and ov(i) >= .5 * (allc[i][2] - allc[i][1])]
            le_r = [r for r, i in enumerate(res) if allc[i][0] == p and ov(i) >= 100]
            if st_r: s5 += 1; rr += 1 / (st_r[0] + 1); s1 += st_r[0] == 0
            if le_r: l5 += 1; l1 += le_r[0] == 0
        n = len(queries); lens = [c[2] - c[1] for c in allc]
        print(f"{name:28} {size:5} {len(allc):7} {st.mean(lens):8.0f} {100*sum(c[4] for c in allc)/len(allc):6.1f}% {s1/n:9.3f} {s5/n:9.3f} {rr/n:6.3f} {l1/n:10.3f} {l5/n:10.3f}")