"""STEP1 (MLX): for several scenarios, compress turn-1 into SP and tokenize a REFERENTIAL turn-2 (needs a
turn-1 value, which therefore lives ONLY in the SP) and a matched CONTROL turn-2 (self-contained — the value
is in the question, so the SP isn't needed). Same SP for both. Save for the per-head attention probe."""
import numpy as np, mlx.core as mx
import sp_mlx

M = sp_mlx.get()
tok, pooler, embT = M["tok"], M["pooler"], M["embT"]
emb = lambda ids: embT(mx.array([ids]))

# (turn-1 fact, referential turn-2 [omits the needed value], control turn-2 [includes it])
SCEN = [
 ("A bakery sells muffins for $4 each. Maria buys 6 muffins. How much does she spend?",
  "I pay with a $50 bill. How much change do I get back?",
  "A toy costs $30 and I pay with a $50 bill. How much change do I get back?"),
 ("Tom reads 15 pages each night for 8 nights. How many pages does he read in total?",
  "If the book has 200 pages, how many pages are left to read?",
  "A book has 200 pages and 120 are already read. How many pages are left to read?"),
 ("A class has 30 students. 40% of them are boys. How many boys are there?",
  "How many of the students are girls?",
  "A class has 30 students and 12 are boys. How many are girls?"),
 ("Jake earns $12 per hour and works 9 hours. How much does he earn?",
  "If he saves half of what he earned, how much does he save?",
  "Jake earned $108. If he saves half of it, how much does he save?"),
 ("Sara has 5 boxes with 24 pencils in each box. How many pencils total?",
  "She gives away 45 pencils. How many does she have left?",
  "Sara has 120 pencils and gives away 45. How many does she have left?"),
]
out = {"n": np.array([len(SCEN)]),
       "bos": np.array([tok.bos_token_id if tok.bos_token_id is not None else tok.encode("")[0]])}
for i, (t1, ref, ctrl) in enumerate(SCEN):
    sp = pooler.forward(emb(tok.encode(t1, add_special_tokens=False)).astype(mx.float32))
    out[f"sp_{i}"] = np.array(sp.astype(mx.float32))[0]
    out[f"ref_{i}"] = np.array(tok.encode(ref, add_special_tokens=False))
    out[f"ctrl_{i}"] = np.array(tok.encode(ctrl, add_special_tokens=False))
np.savez("attn_probe2.npz", **out)
print(f"saved {len(SCEN)} scenarios, SP shape {out['sp_0'].shape}")
print("ATTN_EXPORT2_DONE")