# CUERDAS-021 sealed test generator (new seed). Unseen values: words, capitalized names, multi-token values.
import json, random, itertools, re, hashlib, sys
from transformers import AutoTokenizer
tok = AutoTokenizer.from_pretrained(sys.argv[1]); SEED = 20261017
meta = json.load(open(sys.argv[2])); V = meta['vocabulary']; rows = meta['rows']; SPL = meta['splits']
s020 = json.load(open(sys.argv[3]))
def enc(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}
isid = lambda v: v[:2] in seen_at and v[2:] in seen_sw
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)
dev = set([v for v in uni if isid(v)][:200] + [v for v in uni if not isid(v)][:400])
idx = lambda m: tuple(V[j].index(m[j]) for j in range(4))
s20 = {idx(m) for k in ('sealed_id', 'sealed_ood') for m in s020[k]}
elig = [v for v in itertools.product(*[range(len(x)) for x in V]) if v not in used and v not in dev and v not in s20 and not all(v[j] < (7, 7, 8, 5)[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)]
voc = tok.get_vocab()
words = sorted(t[1:] for t in voc if re.fullmatch(r'Ġ[a-z]{3,14}', t)); names = sorted(t[1:] for t in voc if re.fullmatch(r'Ġ[A-Z][a-z]{2,10}', t))
subs = sorted(t for t in voc if re.fullmatch(r'[a-z]{2,5}', t))
r.shuffle(words); r.shuffle(names); r.shuffle(subs)
old = set(sum(s020['new_words'], [])) | set('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())
taken = {enc(w)[0] for vs in V for w in vs} | {enc(w)[0] for w in old}
NEW = [[] for _ in range(4)]; TYPE = {}
wi = iter(words); ni = iter(names)
def take(it):
    while True:
        w = next(it)
        if len(enc(w)) == 1 and enc(w)[0] not in taken and w not in old: taken.add(enc(w)[0]); return w
for j in range(4):
    for _ in range(3): w = take(wi); NEW[j].append(w); TYPE[w] = 'word'
    for _ in range(3): w = take(ni); NEW[j].append(w); TYPE[w] = 'name'
    a = take(wi); b = take(wi); p = a + ' ' + b; assert len(enc(p)) == 2; NEW[j].append(p); TYPE[p] = 'phrase'
    while True:
        a = take(wi); sb = subs.pop(); pw = a + sb
        if enc(pw) == [enc(a)[0], tok.convert_tokens_to_ids(sb)]: NEW[j].append(pw); TYPE[pw] = 'multitoken_word'; break
excl = sorted({t for j in range(4) for w in NEW[j] for t in enc(w)})
def novel(n_new, n=192):
    out = []
    for k in range(n):
        b = W(sood[k % 400])
        for jj in r.sample(range(4), n_new): b[jj] = r.choice(NEW[jj])
        out.append(b)
    return out
S = {'seed': SEED, 'new_values': NEW, 'types': TYPE, 'exclude_token_ids': excl,
     '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(S, sort_keys=True); open(sys.argv[4], 'w').write(s)
print('sha256', hashlib.sha256(s.encode()).hexdigest(), 'eligible', len(elig))
print(NEW)