"""Method bake-off for detecting 'future-relevant proper/specific information'.

We want the MOST ELEGANT signal that fires on pin-worthy tokens/spans (proper nouns,
codes, prices, dates, numeric values) and stays quiet on generic words. We compare:

  (1) TPC  -- tokenizer fragmentation (tokens-per-char). Form-based, per-token, free.
  (2) BGE  -- the user's hypothesis: do specific spans land in a detectable region of
             BGE-small embedding space? Tested as (a) a linear probe (one 'specificity
             direction'), and (b) distance from the centroid of generic-word embeddings.
  (3) SURP -- a frequency proxy for surprisal: -log unigram prob from the tokenizer's
             merge order (rarer BPE token == higher id-rank == more specific FORM).

Cheap (BGE on CPU + tokenizer only; no LM forward pass). Prints separation numbers so we
can see which signal cleanly splits SPECIFIC from GENERIC -- and, crucially, whether ANY
form-based signal catches '$120' (the case the user says slips through TPC).

Run: python3.12 evals/specificity_probe.py
"""
import sys, os, math
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "runtime"))
import numpy as np

# ---- labeled spans (the thing we care about), with a 'kind' so we can see WHICH types each signal misses
SPECIFIC = [
    ("$120", "money"), ("$500", "money"), ("14C", "code"), ("B3", "code"),
    ("QX7-2291", "code"), ("POL-55821", "code"), ("5588", "number"), ("EMP-90832", "code"),
    ("555-0199", "number"), ("Naruhito", "name"), ("Sapporo", "name"), ("Mochi", "name"),
    ("Helsinki", "name"), ("Apollo", "name"), ("Mei", "name"), ("Hilton", "name"),
    ("hunter2", "code"), ("3pm", "time"), ("Friday", "date"), ("80 cm", "measure"),
    ("O negative", "medical"), ("level B3", "code"), ("bay 12", "code"), ("7 years", "number"),
]
GENERIC = [
    ("the", "stop"), ("planning", "verb"), ("explain", "verb"), ("weekend", "common"),
    ("trip", "common"), ("theme", "common"), ("party", "common"), ("idea", "common"),
    ("help", "verb"), ("something", "common"), ("really", "adv"), ("maybe", "adv"),
    ("about", "prep"), ("really nice", "phrase"), ("a little", "phrase"), ("how are you", "phrase"),
    ("good morning", "phrase"), ("let me", "phrase"), ("i think", "phrase"), ("the meeting", "phrase"),
    ("my budget", "phrase"), ("the total", "phrase"), ("some ideas", "phrase"), ("the cake", "phrase"),
]


def tpc(tok, s):
    ids = tok.encode(s, add_special_tokens=False)
    return len(ids) / max(1, len(s)), len(ids)


def surp_proxy(tok, s):
    # higher BPE token id ~ rarer token (vocab roughly frequency-ordered). Use mean id-rank
    # normalized by vocab size as a crude -log p(form). Robust, no LM needed.
    ids = tok.encode(s, add_special_tokens=False)
    V = tok.vocab_size
    return float(np.mean([i / V for i in ids])) if ids else 0.0


def main():
    from transformers import AutoTokenizer
    from rag import BGERetriever
    here = os.path.dirname(os.path.abspath(__file__))
    tok = AutoTokenizer.from_pretrained(os.path.join(here, "..", "fft_mlx4"))
    bge = BGERetriever()

    spans = [(t, k, 1) for t, k in SPECIFIC] + [(t, k, 0) for t, k in GENERIC]
    texts = [t for t, _, _ in spans]
    y = np.array([lab for _, _, lab in spans])

    # --- signal 1: TPC,  signal 3: surprisal proxy
    feats = {"tpc": [], "surp": []}
    for t, _, _ in spans:
        f, _n = tpc(tok, t); feats["tpc"].append(f)
        feats["surp"].append(surp_proxy(tok, t))
    feats = {k: np.array(v) for k, v in feats.items()}

    # --- signal 2: BGE embeddings of the bare span
    E = bge._encode(texts, is_query=False)                      # (N,384), L2-normed
    gen_centroid = E[y == 0].mean(0); gen_centroid /= np.linalg.norm(gen_centroid)
    bge_dist = 1 - E @ gen_centroid                             # cosine distance from generic centroid

    print("=" * 72)
    print("PER-SPAN signals (sorted by TPC):  spec? | tpc  surp  bgeDist | kind | span")
    order = np.argsort(-feats["tpc"])
    for i in order:
        t, k, lab = spans[i]
        print(f"  {'SPEC' if lab else 'gen '} | {feats['tpc'][i]:.2f} {feats['surp'][i]:.2f} "
              f"{bge_dist[i]:.2f} | {k:8s} | {t}")

    # --- how well does each scalar signal separate spec vs gen?  (AUC = P(spec score > gen score))
    def auc(score):
        ps, ng = score[y == 1], score[y == 0]
        wins = sum((s > g) + 0.5 * (s == g) for s in ps for g in ng)
        return wins / (len(ps) * len(ng))
    print("\nSeparation (AUC, 1.0=perfect):")
    for name, sc in [("tpc", feats["tpc"]), ("surp", feats["surp"]), ("bge_centroid_dist", bge_dist)]:
        print(f"  {name:18s} AUC={auc(sc):.3f}")

    # --- BGE linear probe: is 'specificity' a single linear direction in BGE space?
    from sklearn.linear_model import LogisticRegression
    from sklearn.model_selection import cross_val_predict, StratifiedKFold
    from sklearn.metrics import accuracy_score
    clf = LogisticRegression(max_iter=2000, C=2.0, class_weight="balanced")
    pred = cross_val_predict(clf, E, y, cv=StratifiedKFold(4, shuffle=True, random_state=0))
    print(f"\nBGE linear probe (specificity direction), 4-fold CV acc = {accuracy_score(y, pred):.3f}")
    print("  probe mistakes:")
    for i in range(len(spans)):
        if pred[i] != y[i]:
            print(f"    {'SPEC' if y[i] else 'gen '}->{'SPEC' if pred[i] else 'gen '}: {spans[i][0]}")

    # --- the crucial case: does ANY form signal catch the money value the user flagged?
    print("\nThe '$120 slips through' check:")
    i120 = texts.index("$120")
    print(f"  $120: tpc={feats['tpc'][i120]:.2f} (LOW -> form signal misses it), "
          f"bgeDist={bge_dist[i120]:.2f}, probe={'SPEC' if pred[i120] else 'gen'}")
    print("SPEC_PROBE_DONE")


if __name__ == "__main__":
    main()