"""Matryoshka truncation on the real FastAPI docs (run from experiments/: python embedding-dimension-tradeoffs/bench_dims.py).
Model: mixedbread-ai/mxbai-embed-large-v1 (1024-d, Matryoshka-trained). Docs = h2/h3 sections with >=300 chars body (first 1500 chars, model truncates at 256 tokens).
Queries = the section headings (2+ words) with the model's retrieval prompt. Hit = the section's own doc is retrieved.
Truncate to d dims, re-normalise, exact cosine. Also: int8 scalar quantisation of the 1024-d vectors, and latency/memory of exact search at 100k vectors."""
import re, glob, time, platform, numpy as np, torch, sentence_transformers as st_, faiss
faiss.omp_set_num_threads(2); torch.set_num_threads(2)
from sentence_transformers import SentenceTransformer
H = re.compile(r"^(#{2,3}) (.+?)\s*(\{.*\})?$", re.M)
docs, queries = [], []
for p in sorted(glob.glob('corpus/*.md')):
    t = open(p, encoding='utf-8').read(); ms = list(H.finditer(t))
    for i, m in enumerate(ms):
        end = ms[i+1].start() if i+1 < len(ms) else len(t)
        body = t[m.end():end].strip(); title = m.group(2).strip()
        if len(body) >= 300 and len(title.split()) >= 2:
            docs.append(f"{title}\n{body}"[:1500]); queries.append(title)
print(f"docs={len(docs)} queries={len(queries)}", flush=True)
model = SentenceTransformer('mixedbread-ai/mxbai-embed-large-v1', device='cpu', model_kwargs={'dtype': torch.float32})
model.max_seq_length = 256
t0 = time.time()
D = model.encode([d for d in docs], batch_size=16, normalize_embeddings=False)
Q = model.encode(['Represent this sentence for searching relevant passages: ' + q for q in queries], batch_size=32, normalize_embeddings=False)
print(f"encode seconds={time.time()-t0:.0f}", flush=True)
def nrm(x): return x / np.linalg.norm(x, axis=1, keepdims=True)
def evalq(Dm, Qm):
    S = Qm @ Dm.T; gold = S[np.arange(len(Qm)), np.arange(len(Qm))]
    rank = (S > gold[:, None]).sum(1) + 1
    return (rank == 1).mean(), (rank <= 10).mean(), (1 / rank).mean()
print(f"\n{'dims':>5} {'bytes/vec':>9} {'R@1':>6} {'R@10':>6} {'MRR':>6}")
for d in (1024, 768, 512, 256, 128, 64):
    r1, r10, mrr = evalq(nrm(D[:, :d]), nrm(Q[:, :d]))
    print(f"{d:>5} {d*4:>9} {r1:6.3f} {r10:6.3f} {mrr:6.3f}")
# int8 scalar quantisation of 768-d (per-dimension symmetric scale from docs)
Dn, Qn = nrm(D), nrm(Q); sc = np.abs(Dn).max(0) / 127
Dq = np.round(Dn / sc).astype(np.int8).astype(np.float32) * sc
r1, r10, mrr = evalq(Dq, Qn); print(f"1024 int8 bytes/vec=1024  R@1={r1:.3f} R@10={r10:.3f} MRR={mrr:.3f}")
print("\nExact search, 100k random unit vectors, float32, k=10, median of 100 queries, 2 threads")
print(f"{'dims':>5} {'MB':>7} {'numpy ms':>9} {'faiss ms':>9}")
rng = np.random.default_rng(0); n = 100_000
for d in (256, 768, 3072):
    X = nrm(rng.standard_normal((n, d), dtype=np.float32)); Qs = nrm(rng.standard_normal((100, d), dtype=np.float32))
    tn = []
    for q in Qs:
        s = time.perf_counter(); sc_ = X @ q; i = np.argpartition(-sc_, 10)[:10]; i[np.argsort(-sc_[i])]; tn.append(time.perf_counter() - s)
    ix = faiss.IndexFlatIP(d); ix.add(X); tf = []
    for q in Qs:
        s = time.perf_counter(); ix.search(q[None], 10); tf.append(time.perf_counter() - s)
    print(f"{d:>5} {n*d*4/1e6:7.0f} {np.median(tn)*1e3:9.2f} {np.median(tf)*1e3:9.2f}")
    del X, ix
print(f"\nnumpy {np.__version__}, faiss {faiss.__version__}, sentence-transformers {st_.__version__}, torch {torch.__version__}, python {platform.python_version()}, {platform.processor() or platform.machine()}")