"""Score-threshold calibration for "I do not know" answers. FastAPI docs corpus, heading-aware chunks, MiniLM-L6-v2 + BM25.
30 answerable paraphrase queries (hand-written, from the hybrid-search post) vs 30 unanswerable queries (15 near-domain, 15 unrelated)."""
import re, glob, math, collections, statistics as st, time
import numpy as np
from sentence_transformers import SentenceTransformer

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))]

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"),
]
UNANSWERABLE_NEAR = [  # same domain (Python web APIs), but not covered by the FastAPI docs
    "How do I implement cursor-based pagination for a list endpoint?",
    "How do I rate limit requests per API key?",
    "How do I expose Prometheus metrics from my service?",
    "How do I define a gRPC service and call it from Python?",
    "How do I index documents into Elasticsearch?",
    "How do I provision a server with Terraform?",
    "How do I automate a browser with Selenium?",
    "How do I rotate OAuth refresh tokens?",
    "How do I send email with attachments through SMTP?",
    "How do I schedule a cron job that runs every night?",
    "How do I train a gradient boosting model on tabular data?",
    "How do I cache expensive function results with a TTL?",
    "How do I parse a CSV file with a million rows?",
    "How do I implement full-text search with PostgreSQL tsvector?",
    "How do I build a Discord bot?",
]
UNANSWERABLE_FAR = [  # unrelated to software
    "What is the capital of Australia?",
    "How long should I boil an egg?",
    "Who won the 2018 football world cup?",
    "What is the boiling point of nitrogen?",
    "How do I knit a scarf?",
    "What are the symptoms of the flu?",
    "How many moons does Jupiter have?",
    "Best way to learn to play guitar",
    "What is the plot of Hamlet?",
    "How do I change a flat bicycle tyre?",
    "What year did the Berlin Wall fall?",
    "How do I convert miles to kilometres?",
    "Recommend a good pasta sauce recipe",
    "How do tides work?",
    "What is the tallest mountain in Africa?",
]

ans_q = [q for q, _ in PARAPHRASE]
ans_f = [f for _, f in PARAPHRASE]
all_q = ans_q + UNANSWERABLE_NEAR + UNANSWERABLE_FAR
labels = [1]*len(ans_q) + [0]*(len(UNANSWERABLE_NEAR)+len(UNANSWERABLE_FAR))
kind = ['answerable']*len(ans_q) + ['near']*len(UNANSWERABLE_NEAR) + ['far']*len(UNANSWERABLE_FAR)

# --- score each query: top-1 cosine, top-1 BM25 raw, BM25 matched-query-term fraction ---
def bm25_scores(q):
    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 sc

qe = model.encode(all_q, show_progress_bar=False, convert_to_numpy=True)
qe = qe / np.linalg.norm(qe, axis=1, keepdims=True)
sims = qe @ emb.T
cos_top = sims.max(axis=1)
cos_gap = np.array([np.sort(s)[-1] - np.sort(s)[-2] for s in sims])
bm_top, bm_empty = [], []
for q in all_q:
    sc = bm25_scores(q)
    bm_top.append(max(sc.values()) if sc else 0.0)
    bm_empty.append(not sc)
bm_top = np.array(bm_top)
print(f"\nBM25 returned zero candidates for {sum(bm_empty)} of {len(all_q)} queries "
      f"(answerable: {sum(e for e,l in zip(bm_empty,labels) if l)}, unanswerable: {sum(e for e,l in zip(bm_empty,labels) if not l)})")
# stopwords make BM25 'non-empty' for almost everything; also try a stopword-filtered variant
STOP = set("a an the i do how to what is of in on for and or my me you your it with by from at as be are was can should when who which that this there their".split())
def bm25_scores_nostop(q):
    sc = collections.defaultdict(float)
    for t in set(toks(q)) - STOP:
        for i, f in inv.get(t, ()):
            sc[i] += idf[t] * f * 2.5 / (f + 1.5 * (.25 + .75 * L[i] / avg))
    return sc
bm_ns = np.array([max(bm25_scores_nostop(q).values(), default=0.0) for q in all_q])
print(f"BM25 (stopwords removed) zero candidates: answerable {sum(1 for s,l in zip(bm_ns,labels) if l and s==0)}, "
      f"unanswerable {sum(1 for s,l in zip(bm_ns,labels) if not l and s==0)}")

lab = np.array(labels); kd = np.array(kind)
def summary(name, x):
    print(f"{name:28s} " + "  ".join(
        f"{k}: median {np.median(x[kd==k]):.3f} [min {x[kd==k].min():.3f}, max {x[kd==k].max():.3f}]" for k in ['answerable','near','far']))
print("\nTop-1 score distributions")
summary("cosine top-1", cos_top); summary("BM25 top-1 (raw)", bm_top); summary("BM25 top-1 (no stopwords)", bm_ns)

def auc(x):
    pos, neg = x[lab==1], x[lab==0]
    return (sum((p>n)+0.5*(p==n) for p in pos for n in neg)) / (len(pos)*len(neg))
print("\nAUC (answerable vs unanswerable):")
for name, x in [("cosine top-1", cos_top), ("BM25 top-1 raw", bm_top), ("BM25 no-stopword", bm_ns), ("cosine top1-top2 gap", cos_gap)]:
    print(f"  {name:22s} {auc(x):.3f}")

def sweep(x, thr_list):
    print(f"\n{'thr':>7} | answered-kept | near refused | far refused | all unanswerable refused")
    for t in thr_list:
        keep = x >= t
        ak = keep[lab==1].mean(); nr = (~keep[kd=='near']).mean(); fr = (~keep[kd=='far']).mean(); ur = (~keep[lab==0]).mean()
        print(f"{t:7.2f} | {ak:13.3f} | {nr:12.3f} | {fr:11.3f} | {ur:.3f}")
print("\nCosine threshold sweep (refuse if top-1 cosine < thr)")
sweep(cos_top, [round(0.1*k/2,2) for k in range(0,13)])

# --- calibration / held-out test: even indices calibrate, odd indices test (per class) ---
idx_cal, idx_test = [], []
cnt = collections.Counter()
for i, k in enumerate(kind):
    (idx_cal if cnt[k] % 2 == 0 else idx_test).append(i); cnt[k] += 1
idx_cal, idx_test = np.array(idx_cal), np.array(idx_test)
def best_threshold(x, idx, min_keep):
    # highest-refusal threshold that still answers >= min_keep of answerable calibration queries
    cands = sorted(set(x[idx]))
    best = None
    for t in cands:
        keep = x[idx] >= t
        ak = keep[lab[idx]==1].mean(); ur = (~keep[lab[idx]==0]).mean()
        if ak >= min_keep and (best is None or ur > best[2] or (ur == best[2] and t < best[0])):
            best = (t, ak, ur)
    return best
print("\nCalibrate on half, test on the other half (cosine top-1)")
for mk in [1.0, 0.9]:
    t, ak, ur = best_threshold(cos_top, idx_cal, mk)
    keep = cos_top[idx_test] >= t
    lt = lab[idx_test]
    print(f"  require >= {mk:.0%} answerable kept on calibration -> thr {t:.3f} (cal: kept {ak:.2f}, refused {ur:.2f}); "
          f"TEST: answerable kept {keep[lt==1].mean():.3f} ({keep[lt==1].sum()}/{(lt==1).sum()}), "
          f"unanswerable refused {(~keep[lt==0]).mean():.3f} ({(~keep[lt==0]).sum()}/{(lt==0).sum()})")

print("\nPer-query cosine top-1 (sorted), for the unanswerable queries with highest scores:")
order = np.argsort(-cos_top)
for i in order:
    if lab[i]==0 and cos_top[i] >= np.percentile(cos_top[lab==1], 10):
        print(f"  {cos_top[i]:.3f} [{kind[i]}] {all_q[i]}  -> {chunks[int(sims[i].argmax())][0].split('/')[-1]}")
print("\nLowest-scoring answerable queries:")
for i in np.argsort(cos_top):
    if lab[i]==1 and cos_top[i] <= np.percentile(cos_top[lab==1], 15):
        print(f"  {cos_top[i]:.3f} {all_q[i]}")
