# CUERDAS-020: sealed test generator. Fresh seed, never seen by any earlier iteration.
import json, random, itertools, re, hashlib, sys
from transformers import AutoTokenizer
tok = AutoTokenizer.from_pretrained(sys.argv[1]); SEED = 20261016
meta = json.load(open(sys.argv[2])); V = meta['vocabulary']; rows = meta['rows']; SPL = meta['splits']
def t1(w): return tok.encode(' ' + w, add_special_tokens=False)
used = {tuple(r['values']) for r in rows}
fit = [tuple(rows[i]['values']) for i in SPL['fit']]; seen_at = {v[:2] for v in fit}; seen_sw = {v[2:] for v in fit}
# replicate the 015-019 fresh test (seed 15015) so it is excluded
rng = random.Random(15015); uni = [v for v in itertools.product(*[range(len(x)) for x in V]) if v not in used]; rng.shuffle(uni)
isid = lambda v: v[:2] in seen_at and v[2:] in seen_sw
fresh = set([v for v in uni if isid(v)][:200] + [v for v in uni if not isid(v)][:400])
S013 = (7, 7, 8, 5)
elig = [v for v in itertools.product(*[range(len(x)) for x in V]) if v not in used and v not in fresh and not all(v[j] < S013[j] for j in range(4))]
r = random.Random(SEED); r.shuffle(elig)
sid = [v for v in elig if isid(v)][:200]; sood = [v for v in elig if not isid(v)][:400]
W = lambda v: [V[j][v[j]] for j in range(4)]
# unseen words: single-token, >=10 lowercase letters (the 018 training pool only has 3-9 letters, so these can never have been trained on)
OLD_NEW = 'inspect repair cancel approve archive scan print send sensor battery printer server truck door screen cable broken stable late empty full warm missing urgent night evening weekend noon spring summer winter midnight'.split()
cand = sorted(t[1:] for t in tok.get_vocab() if t.startswith('Ġ') and re.fullmatch(r'[a-z]{10,14}', t[1:]))
r.shuffle(cand); usedtok = {t1(w)[0] for vs in V for w in vs}; NEW = [[] for _ in range(4)]
for w in cand:
    j = min(range(4), key=lambda k: len(NEW[k]))
    if len(NEW[j]) >= 8: break
    if len(t1(w)) == 1 and t1(w)[0] not in usedtok and w not in OLD_NEW: NEW[j].append(w); usedtok.add(t1(w)[0])
def novel(n_new, n=192):
    out = []
    for k in range(n):
        b = W(sood[k % 400]); js = r.sample(range(4), n_new)
        for j in js: b[j] = r.choice(NEW[j])
        out.append(b)
    return out
SEALED = {'seed': SEED, 'n_eligible_universe': len(elig) + len(sid) * 0, 'new_words': NEW,
          'sealed_id': [W(v) for v in sid], 'sealed_ood': [W(v) for v in sood],
          'sealed_1_new': novel(1), 'sealed_2_new': novel(2), 'sealed_4_new': novel(4)}
s = json.dumps(SEALED, sort_keys=True); open(sys.argv[3], 'w').write(s)
print('sha256', hashlib.sha256(s.encode()).hexdigest(), 'eligible', len(elig), {k: len(v) for k, v in SEALED.items() if isinstance(v, list)})