#!/usr/bin/env python3
"""Precomputed word map: everyday words -> library assets by meaning (18 Oct 2026, Part C2).

The Factory has no numpy and no query embedder, so the meaning lookup is done here, at publish time:
every vocabulary entry is embedded with the TEXT model that made the asset DESCRIPTION vectors (gemini-embedding-001,
RETRIEVAL_QUERY for the word, the stored RETRIEVAL_DOCUMENT vector for each asset), cosine against every asset, the best
few kept. The Factory and the Engine just read search/word_map.json: no runtime API call.

Re-runnable: word vectors are cached by text hash under vectors/words/ (only new words are embedded), the output is
byte-stable for the same inputs. Run again after a library publish (new assets) or after editing taxonomy/vocab.txt.

  python tools/build_word_map.py                 # build search/word_map.json
  python tools/build_word_map.py --calibrate     # precision/recall on test/word_map_cases.json for a grid of thresholds
"""
import argparse, hashlib, json, re, sys
from pathlib import Path
import numpy as np

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "lib"))
VEC_SET = "gemini-embedding-2-preview"      # folder of the asset vectors (their `description` array = the TEXT model)
OUT = ROOT / "search" / "word_map.json"
CACHE = ROOT / "vectors" / "words"
VOCAB = ROOT / "taxonomy" / "vocab.txt"
CASES = ROOT / "test" / "word_map_cases.json"
THRESHOLD = 0.68      # calibrated on test/word_map_cases.json: all 12 negatives clean, recall 0.67, precision 0.76
MARGIN = 0.04       # keep only assets within 0.04 of the word's best score
TEMPLATE = "a 3D model of {}"   # a bare word sits far from a scene description; this puts it in the same register
PER_WORD = 8
BATCH = 20


def key(a):
    return hashlib.sha256((a["id"] + "@" + a["release"]).encode()).hexdigest()


def norm(w):
    return re.sub(r"\s+", " ", re.sub(r"[^a-z0-9]+", " ", w.lower())).strip()


def tokens_of(w):
    """Same stemming as library_search.tokens, loaded from the Factory file so both sides agree."""
    sys.path.insert(0, str(ROOT / "omg" / "factory"))
    import library_search
    return library_search.tokens(w)


def vocabulary(index):
    words = set()
    for kind, node in index["tree"].items():
        for g in node["groups"].values():
            words.update(norm(c) for c in g["classes"])
    words.update(norm(a) for a in index.get("aliases", {}))
    for a in index["assets"]:
        for t in re.split(r"[^a-z]+", (a.get("name") or "").lower() + " " + " ".join(a.get("tags") or [])):
            if len(t) >= 3:
                words.add(t)
    if VOCAB.exists():
        words.update(norm(x) for x in VOCAB.read_text().splitlines())
    return sorted(w for w in words if len(w) >= 3)


def word_vectors(words, calls):
    """text -> unit vector, cached in vectors/words/<2 hex>.npz shards keyed by sha256(model|text)."""
    import embeddings as e
    CACHE.mkdir(parents=True, exist_ok=True)
    shards, out, todo = {}, {}, []
    for w in words:
        h = hashlib.sha256((e.TEXT + "|query|" + TEMPLATE + "|" + w).encode()).hexdigest()
        sh = shards.get(h[:2])
        if sh is None:
            p = CACHE / (h[:2] + ".npz")
            sh = shards[h[:2]] = {k: v for k, v in np.load(p).items()} if p.exists() else {}
        if h in sh:
            out[w] = sh[h].astype(np.float32)
        else:
            todo.append((w, h))
    for i in range(0, len(todo), BATCH):
        chunk = todo[i:i + BATCH]
        preds = e.predict(e.TEXT, [{"content": TEMPLATE.format(w), "task_type": "RETRIEVAL_QUERY"} for w, _ in chunk], {"outputDimensionality": 768})
        calls[0] += 1
        for (w, h), p in zip(chunk, preds):
            v = e.unit(p["embeddings"]["values"])
            shards[h[:2]][h] = v.astype("<f2")
            out[w] = v
    if todo:
        for s, d in shards.items():
            np.savez_compressed(CACHE / (s + ".npz"), **{k: d[k] for k in sorted(d)})
    return out


def score_table(index, words, calls):
    d = ROOT / "vectors" / VEC_SET
    assets = [a for a in sorted(index["assets"], key=lambda a: a["id"]) if (d / (key(a) + ".npz")).exists()]
    mat = np.array([e_unit(np.load(d / (key(a) + ".npz"))["description"]) for a in assets])
    wv = word_vectors(words, calls)
    W = np.array([wv[w] for w in words])
    return assets, W @ mat.T


def e_unit(v):
    v = np.asarray(v, dtype=np.float32)
    return v / max(float(np.linalg.norm(v)), 1e-12)


def pick(row, assets, threshold, margin):
    order = sorted(range(len(assets)), key=lambda j: (-round(float(row[j]), 3), assets[j]["id"]))[:PER_WORD]
    top = round(float(row[order[0]]), 3)
    return [[assets[j]["id"], round(float(row[j]), 3)] for j in order
            if round(float(row[j]), 3) >= threshold and round(float(row[j]), 3) >= top - margin]


def calibrate(index, calls):
    cases = json.loads(CASES.read_text())
    words = sorted({c["word"] for c in cases["positives"] + cases["negatives"]})
    assets, S = score_table(index, [norm(w) for w in words], calls)
    row = {w: S[i] for i, w in enumerate(norm(w) for w in words)}
    print(f"{'thr':>5} {'margin':>6} | recall (positives hit) | negatives clean | precision (hits on expected / hits)")
    best = None
    for margin in (0.02, 0.04, 0.06, 1.0):
        for thr in [round(0.50 + 0.01 * i, 2) for i in range(0, 30)]:
            hit = tot = good = 0
            for c in cases["positives"]:
                h = pick(row[norm(c["word"])], assets, thr, margin)
                tot += len(h)
                g = sum(1 for i, _ in h if i.startswith(tuple(c["expect"])))
                good += g
                hit += g > 0
            clean = sum(1 for c in cases["negatives"] if not pick(row[norm(c["word"])], assets, thr, margin))
            rec, neg = hit / len(cases["positives"]), clean / len(cases["negatives"])
            prec = good / tot if tot else 0.0
            score = (neg == 1.0, 2 * rec * prec / (rec + prec) if rec + prec else 0.0)
            if best is None or score > best[0]:
                best = (score, thr, margin, rec, neg, prec)
            if margin == 0.04 and (thr * 100) % 2 == 0:
                print(f"{thr:5.2f} {margin:6.2f} | {rec:6.2f} | {neg:6.2f} | {prec:6.2f}")
    print("BEST (all negatives clean, then best F1): thr=%.2f margin=%.2f recall=%.2f negatives_clean=%.2f precision=%.2f" % best[1:])
    for c in cases["positives"] + cases["negatives"]:
        h = pick(row[norm(c["word"])], assets, best[1], best[2])
        print(f"  {c['word']:14s} {'+' if 'expect' in c else '-'} {[(i.split('/')[-1], s) for i, s in h][:4]}")
    return best


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--calibrate", action="store_true")
    ap.add_argument("--threshold", type=float, default=THRESHOLD)
    ap.add_argument("--margin", type=float, default=MARGIN)
    ap.add_argument("--index", default=str(ROOT / "index.json"))
    ap.add_argument("--out", default=str(OUT))
    args = ap.parse_args()
    index = json.loads(Path(args.index).read_text())
    calls = [0]
    if args.calibrate:
        calibrate(index, calls)
        print("embedding calls:", calls[0])
        return
    words = vocabulary(index)
    assets, S = score_table(index, words, calls)
    table, seen = {}, set()
    for i, w in enumerate(words):
        stem = " ".join(tokens_of(w))
        if not stem or stem in seen:      # one entry per stemmed form (first alphabetical wins: deterministic)
            continue
        seen.add(stem)
        hits = pick(S[i], assets, args.threshold, args.margin)
        if hits:
            table[stem] = hits
    doc = {"model": "gemini-embedding-001", "threshold": args.threshold, "margin": args.margin,
           "built_from_index": index.get("built"), "assets": len(assets), "words": dict(sorted(table.items()))}
    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(doc, sort_keys=True, separators=(",", ":")) + "\n")
    print(f"vocabulary {len(words)}, words with a hit {len(table)}, embedding calls {calls[0]}, wrote {out} ({out.stat().st_size // 1024} KB)")


if __name__ == "__main__":
    main()