"""Chunk overlap sweep (run from experiments/: python chunk-overlap-does-it-help/bench_overlap.py).
Fixed windows of SIZE characters with overlap 0/10/20/40% of SIZE; BM25 top-5 on FastAPI docs heading queries.
Strict/lenient hit rules as in the chunk size sweep post.
"""
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, pct):
    step = size - size * pct // 100; out, i = [], 0
    while i < len(text):
        out.append((i, min(len(text), i + size), ''))
        if out[-1][1] == len(text): break
        i += step
    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 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"{'size':>5} {'overlap':>8} {'chunks':>7} {'index x':>8} {'strict@1':>9} {'strict@5':>9} {'MRR':>6} {'lenient@1':>10} {'lenient@5':>10}")
for size in [200, 400, 800, 1200]:
    for pct in [0, 10, 20, 40]:
        allc = []
        for p, t in docs.items():
            for a, b, _ in fixed(t, size, pct): allc.append((p, a, b))
        search = bm25_index([toks(docs[p][a:b]) for p, a, b 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); tot = sum(b - a for _, a, b in allc)
        print(f"{size:5} {pct:7}% {len(allc):7} {tot/sum(map(len, docs.values())):8.2f} {s1/n:9.3f} {s5/n:9.3f} {rr/n:6.3f} {l1/n:10.3f} {l5/n:10.3f}")