# CUERDAS-021 common: sealed set, varied replacement pool, training loop, evaluation.
import os, sys, json, hashlib, random, re, time
import torch, torch.nn as nn, torch.nn.functional as F
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import run018 as R
INP = os.environ['INP']; OD = os.environ['OD']
raw = open(f'{INP}/sealed021.json').read(); H = hashlib.sha256(raw.encode()).hexdigest()
assert H == 'a86a6a5ee375ae67a3e79f969c9553b0e5548ed3b0e80e5315862e6552ae72bc', H
S = json.loads(raw); NEW = S['new_values']; EXCL = set(S['exclude_token_ids'])
SETS = ['sealed_id', 'sealed_ood', 'sealed_1_new', 'sealed_2_new', 'sealed_4_new']
SMOKE = bool(os.environ.get('SMOKE'))
if SMOKE:  # dry run on dev messages only; sealed sets are never touched
    for s_ in SETS: S[s_] = R.FRESH['ood'][:6]
STEPS = 4 if SMOKE else 3000; EVERY = 2 if SMOKE else 250
t1 = R.t1; tok = R.tok; Emb = R.Emb; D = R.D; FIELDS = R.FIELDS; V = R.V
CAND = [torch.tensor(sorted({t1(w)[0] for w in V[j] + NEW[j]})) for j in range(4)]
assert all(len(CAND[j]) == len(V[j]) + 8 for j in range(4))
def gidx(cand, j, w): return int((cand[j] == t1(w)[0]).nonzero()[0])
# ---- varied pool (training) and held-out validation pool; all disjoint from sealed tokens and from 014/dev/020 values ----
voc = tok.get_vocab(); BAN = EXCL | R.BANNED
def ok_tok(t): return voc[t] not in BAN
WORDS = sorted(t[1:] for t in voc if re.fullmatch(r'Ġ[a-z]{3,14}', t) and ok_tok(t))
NAMES = sorted(t[1:] for t in voc if re.fullmatch(r'Ġ[A-Z][a-z]{2,10}', t) and ok_tok(t))
SUBS = sorted(t for t in voc if re.fullmatch(r'[a-z]{2,5}', t))
g0 = random.Random(2101); g0.shuffle(WORDS); g0.shuffle(NAMES)
VW, VN = WORDS[-300:], NAMES[-300:]; WORDS, NAMES = WORDS[:-300], NAMES[:-300]
def sample_value(g, W_=None, N_=None):
    W_ = W_ or WORDS; N_ = N_ or NAMES; k = g.randrange(5)
    if k == 0: return g.choice(W_)
    if k == 1: return g.choice(N_)
    if k == 2: return str(g.randrange(10000))
    if k == 3: return ' '.join(g.choice(W_) for _ in range(g.choice((2, 3))))
    return g.choice(W_) + g.choice(SUBS)
def sample_msgs(g, bs=12):
    out = []
    for _ in range(bs):
        m = list(g.choice(R.FIT))
        for j in range(4):
            if g.random() < .4: m[j] = sample_value(g)
        out.append(m)
    return out
# restricted transfer validation: 50 held-out values per field (no numbers: first token not unique), options = field vocab + those
gv = random.Random(2102); CANDV = []; VPOOL = []
for j in range(4):
    used = {t1(w)[0] for w in V[j]}; ws = []
    while len(ws) < 50:
        k = gv.randrange(3); w = gv.choice(VW) if k == 0 else gv.choice(VN) if k == 1 else gv.choice(VW) + ' ' + gv.choice(VW)
        if t1(w)[0] not in used: used.add(t1(w)[0]); ws.append(w)
    VPOOL.append(ws); CANDV.append(torch.tensor(sorted(used)))
VAL_TR = [[gv.choice(VPOOL[j]) for j in range(4)] for _ in range(96)]
VALC = R.VAL[:96]
if SMOKE: VALC, VAL_TR = VALC[:6], VAL_TR[:6]
# ---- generic readout: rows = prefix (embeddings) + body tokens ----
def run_h(embs):
    L = max(e.shape[0] for e in embs); B = len(embs); x = torch.zeros(B, L, D); m = torch.zeros(B, L, dtype=torch.long)
    for i, e in enumerate(embs): x[i, :e.shape[0]] = e; m[i, :e.shape[0]] = 1
    h = R.lm.model(inputs_embeds=x, attention_mask=m).last_hidden_state[torch.arange(B), m.sum(1) - 1]
    return h, R.lm.lm_head(h)
def k2_rows(cod, msgs, tpl):
    Sl = cod(torch.stack([torch.stack([R.lex(w) for w in m]) for m in msgs]))
    return [torch.cat([Sl[i], Emb[R.sufids(tpl, j)]]) for i in range(len(msgs)) for j in range(4)]
def text_rows(msgs, tpl, fmt, prefix=None):
    rows = []
    for m in msgs:
        for j in range(4):
            e = Emb[tok.encode(fmt.format(**dict(zip(FIELDS, m))) + tpl.format(field=FIELDS[j]), add_special_tokens=False)]
            rows.append(e if prefix is None else torch.cat([prefix, e]))
    return rows
@torch.no_grad()
def per_ex(rowfn, msgs, cand, bs=24):
    out = []
    for b in range(0, len(msgs), bs):
        mb = msgs[b:b+bs]; _, lg = run_h(rowfn(mb)); good = [True] * len(mb)
        for r_ in range(len(mb) * 4):
            i, j = divmod(r_, 4); good[i] &= int(lg[r_][cand[j]].argmax()) == gidx(cand, j, mb[i][j])
        out += good
    return [float(x) for x in out]
def val_score(rowfn_main):
    return sum(per_ex(rowfn_main, VALC, R.CAND)) / len(VALC) + sum(per_ex(rowfn_main, VAL_TR, CANDV)) / len(VAL_TR)
def train(params, rowfn, teach, seed, tag, log, steps=None, lr=5e-4, every=None, set_mode=lambda t: None):
    """Same loss/schedule/selection as 018 distill2: CE + KL(teacher) + 0.5(1-cos) ; OneCycle; select on validation only."""
    steps = steps or STEPS; every = every or EVERY
    g = random.Random(seed + 2100); opt = torch.optim.AdamW(params, lr=lr, weight_decay=1e-3)
    sch = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=lr, total_steps=steps, pct_start=.05)
    def snap(): return [p.detach().clone() for p in params]
    set_mode(False); v = val_score(lambda mb: rowfn(mb, R.SUF_MAIN)); set_mode(True); best = (v, snap(), 0); log({'run': tag, 'step': 0, 'val': v})
    t0 = time.time()
    for s_ in range(1, steps + 1):
        msgs = sample_msgs(g); tpl = g.choice(R.SUF_TRAIN)
        with torch.no_grad(): ht, lt = R.readout_h(teach, msgs, tpl)
        hs, ls = run_h(rowfn(msgs, tpl)); tg = torch.tensor([t1(m[j])[0] for m in msgs for j in range(4)])
        loss = F.cross_entropy(ls, tg) + F.kl_div(F.log_softmax(ls, -1), F.log_softmax(lt, -1), log_target=True, reduction='batchmean') + .5 * (1 - F.cosine_similarity(hs, ht, -1).mean())
        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(params, 1.0); opt.step(); sch.step()
        if s_ % every == 0:
            set_mode(False); v = val_score(lambda mb: rowfn(mb, R.SUF_MAIN)); set_mode(True)
            log({'run': tag, 'step': s_, 'val': v, 'loss': float(loss), 's': round(time.time() - t0)})
            if v > best[0]: best = (v, snap(), s_)
    with torch.no_grad():
        for p, b in zip(params, best[1]): p.copy_(b)
    set_mode(False); return best[2], best[0]
def load_teacher(sd):
    t = R.Binder(sd); t.load_state_dict(torch.load(f'{INP}/binder4_1000_s{sd}.pt', map_location=R.DEV)); return t.eval().requires_grad_(False)
def opener(name):
    f = open(f'{OD}/{name}', 'w')
    def log(d): f.write(json.dumps(d) + '\n'); f.flush(); print(json.dumps(d), flush=True)
    return log