"""RRF k sweep: what the constant k in 1/(k+rank) does on real BM25 + vector rankings.
Setup (corpus, chunker, BM25, MiniLM embeddings, 3 query sets) is identical to the hybrid-search post.
Fusion here is plain, equal-weight RRF: score(d) = sum over retrievers of 1/(k+rank), rank starting at 1;
a document absent from a retriever's top-50 contributes 0 for that retriever.
"""
import re, glob, math, collections, statistics as st, time
import numpy as np
from sentence_transformers import SentenceTransformer

K_RRF = 60
CAND_DEPTH = 50  # how deep each retriever's ranking feeds the fusion

TOK = re.compile(r"[a-z0-9_]+")
def toks(s): return TOK.findall(s.lower())

HEADING = re.compile(r"^(#{1,3}) (.+?)\s*(\{.*\})?$", re.M)
def sections(text):
    ms = list(HEADING.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

SIZE = 800
def paragraphs(text):
    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_chunks(text):
    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]):
            out.append((s + a, s + b, path))
    return out

docs = {p: open(p, encoding='utf-8').read() for p in sorted(glob.glob('../corpus/*.md'))}

chunks = []  # (doc, start, end, path, text, embed_text)
for p, t in docs.items():
    for a, b, path in heading_aware_chunks(t):
        body = t[a:b]
        embed = f"{path}\n{body}" if path else body
        chunks.append((p, a, b, path, body, embed))
print(f"docs={len(docs)} chars={sum(map(len, docs.values())):,} chunks={len(chunks)}")

# --- BM25 index ---
tf = [collections.Counter(toks(c[5])) for c in chunks]
N = len(chunks)
df = collections.Counter(t for c in tf for t in c)
L = [sum(c.values()) for c in tf]
avg = sum(L) / 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))

def bm25_search(q, k=CAND_DEPTH):
    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 [i for i, _ in sorted(sc.items(), key=lambda x: -x[1])[:k]]

# --- vector index ---
t0 = time.time()
model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
print(f"model loaded in {time.time()-t0:.1f}s")
t0 = time.time()
emb = model.encode([c[5] for c in chunks], batch_size=64, show_progress_bar=False, convert_to_numpy=True)
emb = emb / np.linalg.norm(emb, axis=1, keepdims=True)
print(f"embedded {len(chunks)} chunks in {time.time()-t0:.1f}s")

def vec_search_batch(queries, k=CAND_DEPTH):
    q = model.encode(queries, show_progress_bar=False, convert_to_numpy=True)
    q = q / np.linalg.norm(q, axis=1, keepdims=True)
    sims = q @ emb.T
    return [list(np.argsort(-sims[i])[:k]) for i in range(len(queries))]

# --- heading queries (paraphrase-ish) ---
heading_queries = []
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:
            heading_queries.append((p, s, e, title))
print(f"heading queries: {len(heading_queries)}")

def heading_hit(i, p, s, e):
    cp, ca, cb, *_ = chunks[i]
    if cp != p: return False
    return (min(e, cb) - max(s, ca)) >= 0.5 * (cb - ca)

# --- identifier queries (exact-token) ---
CODE = re.compile(r"`([A-Za-z_][A-Za-z0-9_.]{3,})`")
occ = collections.defaultdict(set)  # identifier -> set of chunk indices
for ci, c in enumerate(chunks):
    for m in CODE.finditer(c[4]):
        ident = m.group(1)
        if ('_' in ident or '.' in ident or ident.lower() != ident) :
            occ[ident].add(ci)
identifier_queries = [(ident, next(iter(idxs))) for ident, idxs in occ.items() if len(idxs) == 1]
print(f"identifier queries: {len(identifier_queries)}")

# --- paraphrase queries (hand-written, no verbatim overlap with the target section's heading) ---
# Each maps to one corpus file that is narrowly scoped to that topic. A hit = any chunk from that file.
PARAPHRASE = [
    ("How do I run code after sending the response to the client?", "tutorial_background-tasks.md"),
    ("How do I change the HTTP status code an endpoint returns without raising an exception?", "advanced_response-change-status-code.md"),
    ("How do I add custom headers to an HTTP response?", "advanced_response-headers.md"),
    ("How do I set a cookie on the response?", "advanced_response-cookies.md"),
    ("How do I accept a persistent two-way connection from a browser?", "advanced_websockets.md"),
    ("How do I run my API behind a reverse proxy that adds a path prefix?", "advanced_behind-a-proxy.md"),
    ("How can I load configuration from environment variables?", "advanced_settings.md"),
    ("How do I generate a TypeScript or Python client from my API automatically?", "advanced_generate-clients.md"),
    ("How can I mount one app inside another at a given path?", "advanced_sub-applications.md"),
    ("How do I run startup and shutdown code for my application?", "advanced_events.md"),
    ("How can I add logic that runs for every request, like a timing header?", "tutorial_middleware.md"),
    ("How do I run my app with multiple worker processes?", "deployment_server-workers.md"),
    ("How do I package my app into a container image?", "deployment_docker.md"),
    ("How do I serve my API over an encrypted connection?", "deployment_https.md"),
    ("How do I read extra parameters from the URL's query string?", "tutorial_query-params.md"),
    ("How do I capture a variable part of the URL path?", "tutorial_path-params.md"),
    ("How can a client upload a file to my API?", "tutorial_request-files.md"),
    ("How do I accept both form fields and an uploaded file in one request?", "tutorial_request-forms-and-files.md"),
    ("How do I return a custom error message and status code when something goes wrong?", "tutorial_handling-errors.md"),
    ("How can I restrict which fields are included in the JSON I return?", "tutorial_response-model.md"),
    ("How do I read a value sent in the request headers?", "tutorial_header-params.md"),
    ("How do I connect my API to a relational database?", "tutorial_sql-databases.md"),
    ("How can I organize a large API into multiple files with shared routing?", "tutorial_bigger-applications.md"),
    ("How do I add a summary and description to an endpoint for the interactive docs?", "tutorial_path-operation-configuration.md"),
    ("How can dependencies be shared and reused across many endpoints?", "tutorial_dependencies_index.md"),
    ("How do I let a browser on a different origin call my API?", "tutorial_cors.md"),
    ("How do I validate the shape of a JSON request body against a schema?", "tutorial_body.md"),
    ("How do I return different response types, like a file download?", "advanced_custom-response.md"),
    ("How can I require a logged-in user before an endpoint runs?", "tutorial_security_first-steps.md"),
    ("How can I make a numeric URL parameter require a minimum or maximum value?", "tutorial_path-params-numeric-validations.md"),
]
import os
paraphrase_queries = [(q, f) for q, f in PARAPHRASE]
print(f"paraphrase queries: {len(paraphrase_queries)}")


import os
paraphrase_queries = [(q, f) for q, f in PARAPHRASE]
h_bm25_ranks = [bm25_search(t) for _, _, _, t in heading_queries]
h_vec_ranks = vec_search_batch([t for _, _, _, t in heading_queries])
i_bm25_ranks = [bm25_search(q) for q, _ in identifier_queries]
i_vec_ranks = vec_search_batch([q for q, _ in identifier_queries])
p_bm25_ranks = [bm25_search(q) for q, _ in paraphrase_queries]
p_vec_ranks = vec_search_batch([q for q, _ in paraphrase_queries])

def p_hit(i, qi): return os.path.basename(chunks[i][0]) == paraphrase_queries[qi][1]
def h_hit(i, qi):
    p, s, e, _ = heading_queries[qi]; return heading_hit(i, p, s, e)
def i_hit(i, qi): return i == identifier_queries[qi][1]

sets = {
  "heading": (h_bm25_ranks, h_vec_ranks, h_hit),
  "ident": (i_bm25_ranks, i_vec_ranks, i_hit),
  "para": (p_bm25_ranks, p_vec_ranks, p_hit),
}

def fuse(bm, ve, k):
    sc = collections.defaultdict(float)
    for r, i in enumerate(bm): sc[i] += 1 / (k + r + 1)
    for r, i in enumerate(ve): sc[i] += 1 / (k + r + 1)
    return [i for i, _ in sorted(sc.items(), key=lambda x: (-x[1], x[0]))]

def mrr5(name, k):
    bms, ves, hit = sets[name]; rr = 0
    for qi, (bm, ve) in enumerate(zip(bms, ves)):
        top = fuse(bm, ve, k)[:5]
        for r, i in enumerate(top):
            if hit(i, qi): rr += 1 / (r + 1); break
    return rr / len(bms)

def mrr5_single(name, which):
    bms, ves, hit = sets[name]; rr = 0
    for qi, rk in enumerate(bms if which == "bm25" else ves):
        for r, i in enumerate(rk[:5]):
            if hit(i, qi): rr += 1 / (r + 1); break
    return rr / len(bms)

print("\n== single retrievers, MRR@5 ==")
for n in sets: print(f"{n:8s} bm25={mrr5_single(n,'bm25'):.3f} vector={mrr5_single(n,'vec'):.3f}")

print("\n== rank weights 1/(k+rank): weight of rank 1 divided by weight of rank r ==")
print(f"{'k':>6} | {'r=2':>6} {'r=5':>6} {'r=10':>6} {'r=50':>6}")
KS = [0, 1, 5, 10, 20, 40, 60, 100, 200, 1000]
for k in KS:
    w = lambda r: 1 / (k + r)
    print(f"{k:6d} | {w(1)/w(2):6.2f} {w(1)/w(5):6.2f} {w(1)/w(10):6.2f} {w(1)/w(50):6.2f}")

print("\n== fused MRR@5 by k (equal weights) ==")
print(f"{'k':>6} | {'heading':>8} {'ident':>8} {'para':>8} | {'mean':>6}")
res = {}
for k in KS:
    v = [mrr5(n, k) for n in sets]; res[k] = v
    print(f"{k:6d} | {v[0]:8.3f} {v[1]:8.3f} {v[2]:8.3f} | {sum(v)/3:6.3f}")
best = max(KS, key=lambda k: sum(res[k])); worst = min(KS, key=lambda k: sum(res[k]))
print(f"best k by mean MRR: {best} ({sum(res[best])/3:.3f}); worst: {worst} ({sum(res[worst])/3:.3f}); k=60: {sum(res[60])/3:.3f}")
print(f"spread (max-min of per-set MRR across all k): " + ", ".join(f"{n}={max(res[k][j] for k in KS)-min(res[k][j] for k in KS):.3f}" for j, n in enumerate(sets)))

print("\n== how often the fused top-1 equals each retriever's top-1 (all queries pooled) ==")
print(f"{'k':>6} | {'=bm25 top1':>10} {'=vec top1':>10} {'either':>8}")
for k in KS:
    nb = nv = ne = n = 0
    for name, (bms, ves, _) in sets.items():
        for bm, ve in zip(bms, ves):
            f = fuse(bm, ve, k)[0]; n += 1
            a = f == bm[0]; b = f == ve[0]; nb += a; nv += b; ne += a or b
    print(f"{k:6d} | {nb/n:10.3f} {nv/n:10.3f} {ne/n:8.3f}")

# Tiny worked example of the k effect: doc A is rank 1 in one list, absent in the other;
# doc B is rank 5 in both. Which wins as k grows?
print("\n== worked example: A = rank 1 in one list only; B = rank 5 in both lists ==")
for k in KS:
    a = 1 / (k + 1); b = 2 / (k + 5)
    print(f"k={k:5d}: A={a:.5f} B={b:.5f} -> {'A' if a > b else 'B'} wins")