"""teach.py -- a quick learned guess of which piece it will pick, from the turns it has already played (runs/*/cands.jsonl).
Ivy, Oct 23 2026: "we really need to attach a nn to every choice it makes so it can make snap decisions when very certain ... It should
be reasonably sure which pieces it's likely to choose so it can pretend to weigh them showing it deliberate and slow down."
Features per candidate: every number in its measurements, plus how well it matches the wish (asked-for kind, asked-for colour host,
left/right/up words in the wish vs where it sits, aimed or surprise). Target: Jev's probability. Tested by leaving whole builds out."""
import json, glob, os, re, sys, pickle
os.environ['HIP_VISIBLE_DEVICES'] = ''; os.environ['CUDA_VISIBLE_DEVICES'] = ''   # torch on Poppy's iGPU ABORTS the process (rocBLAS, gfx900)
import numpy as np
HERE = os.path.dirname(os.path.abspath(__file__)); RUNS = os.path.join(HERE, '..', 'runs')
sys.path.insert(0, os.path.join(HERE, '..'))
KINDS = None
def kinds():
    global KINDS
    if KINDS is None:
        import importlib.util
        src = open(os.path.join(HERE, '..', 'imagine.py')).read(); m = re.search(r"KINDS = (\{.*?\})\n", src, re.S); KINDS = eval(m.group(1))
    return KINDS
def flat(d, pre=''):
    out = {}
    for k, v in (d or {}).items():
        if isinstance(v, dict): out.update(flat(v, pre + k[:24] + '/'))
        elif isinstance(v, bool): out[pre + k[:40]] = float(v)
        elif isinstance(v, (int, float)): out[pre + k[:40]] = float(v)
    return out
_EMB = {}; _MODEL = None
import threading; _LOAD = threading.Lock()   # the server's warm-up and a game thread loaded the model at once and broke the import
def emb(texts):
    """MiniLM sentence embeddings (the same local model the idea engine's taste net uses), cached"""
    global _MODEL
    need = [t for t in texts if t not in _EMB]
    if need:
      with _LOAD:
        if _MODEL is None:
            from sentence_transformers import SentenceTransformer; _MODEL = SentenceTransformer('all-MiniLM-L6-v2', device='cpu')   # Poppy's iGPU can't run torch (rocBLAS has no kernels for gfx900)
        for t, v in zip(need, _MODEL.encode(need, normalize_embeddings=True, batch_size=64)): _EMB[t] = v
    return [_EMB[t] for t in texts]
def cand_text(c):
    d = c.get('desc') or {}; w = d.get('where') or {}
    lr = next((v for k, v in w.items() if k.startswith('left or right')), 0) or 0; ht = next((v for k, v in w.items() if k.startswith('height')), 0) or 0
    return (f"{d.get('description', '')} It joins the {w.get('joins the', '')}, "
            f"{'on the left' if lr < -0.33 else 'on the right' if lr > 0.33 else 'in the middle'}, {'up high' if ht > 0.66 else 'down low' if ht < 0.33 else 'halfway up'}.")
def features(turn, c):
    d = c.get('desc') or {}; f = flat({k: v for k, v in d.items() if k != 'description'})
    wish = (turn.get('next') or '').lower(); want_kind = str(turn.get('want_piece') or '').lower(); want_on = str(turn.get('want_on') or '').lower()
    w = d.get('where') or {}
    lr = next((v for k, v in w.items() if k.startswith('left or right')), 0) or 0; ht = next((v for k, v in w.items() if k.startswith('height')), 0) or 0
    f['aimed'] = float(bool(c.get('aimed', True)))
    ct, wt, mt = emb([cand_text(c), turn.get('next') or '', turn.get('might_be') or ''])
    f['sim_wish'] = float(ct @ wt); f['sim_idea'] = float(ct @ mt)
    f['kind_match'] = float(c.get('name') in (kinds().get(want_kind) or []))
    f['colour_match'] = float(want_on and want_on in str(w.get('joins the', '')).lower())
    f['wish_left'] = float('left' in wish) * (-lr); f['wish_right'] = float('right' in wish) * lr
    f['wish_up'] = float(any(x in wish for x in ('top', 'up', 'high', 'tall'))) * ht; f['wish_low'] = float(any(x in wish for x in ('under', 'low', 'bottom', 'below'))) * (1 - ht)
    f['wish_long'] = float('long' in wish) * float(c.get('name') in (kinds().get('long') or []))
    f['wish_small'] = float(any(x in wish for x in ('small', 'tiny', 'little'))) * float(c.get('name') in (kinds().get('small') or []))
    return f
def load():
    T = []
    for fn in sorted(glob.glob(os.path.join(RUNS, '*', 'cands.jsonl'))):
        run = os.path.basename(os.path.dirname(fn))
        try:                                     # a stopped game is not learned from until it's written down (Ivy, Oct 23 2026)
            if json.load(open(os.path.join(os.path.dirname(fn), 'state.json'))).get('unwritten'): continue
        except (OSError, ValueError): pass
        for l in open(fn):
            t = json.loads(l); cs = [c for c in t['candidates'] if c.get('p') is not None]
            if t.get('forced'): continue           # nothing was right; the engine picked for it (Oct 23 2026) -- not its choice
            if len(cs) >= 2: T.append((run, t, cs))
    return T
def matrix(T, cols=None):
    emb([cand_text(c) for _, t, cs in T for c in cs] + [t.get('next') or '' for _, t, _ in T] + [t.get('might_be') or '' for _, t, _ in T])
    F = [[features(t, c) for c in cs] for _, t, cs in T]
    if cols is None: cols = sorted({k for turn in F for f in turn for k in f})
    X = [np.array([[f.get(k, 0.0) for k in cols] for f in turn]) for turn in F]
    return X, cols
def fit(Xs, ys):
    from sklearn.ensemble import HistGradientBoostingRegressor
    m = HistGradientBoostingRegressor(max_iter=200, learning_rate=0.06, max_leaf_nodes=15, min_samples_leaf=10)
    m.fit(np.vstack(Xs), np.concatenate(ys)); return m
def softmax(z, temp=0.08): z = np.asarray(z) / temp; z = z - z.max(); e = np.exp(z); return e / e.sum()
if __name__ == '__main__':
    T = load(); X, cols = matrix(T); y = [np.array([c['p'] for c in cs]) for _, _, cs in T]
    runs = sorted({r for r, _, _ in T}); folds = [runs[i::5] for i in range(5)]
    hits, chance, conf_hits = 0, 0.0, []
    for fold in folds:
        tr = [i for i, (r, _, _) in enumerate(T) if r not in fold]; te = [i for i, (r, _, _) in enumerate(T) if r in fold]
        m = fit([X[i] for i in tr], [y[i] for i in tr])
        for i in te:
            pr = softmax(m.predict(X[i])); g = int(np.argmax(pr)); truth = int(np.argmax(y[i]))
            hits += g == truth; chance += 1 / len(y[i]); conf_hits.append((float(pr.max()), g == truth))
    n = len(T); print(f'{n} turns from {len(runs)} builds, {len(cols)} features')
    print(f'left-out builds: guessed its pick {hits}/{n} = {hits / n:.0%} (chance {chance / n:.0%})')
    for th in (0.5, 0.7, 0.85):
        sel = [h for p, h in conf_hits if p >= th]
        print(f'  when its guess is >= {th:.0%} sure: {len(sel)} turns ({len(sel) / n:.0%}), right {sum(sel)}/{len(sel)} = {sum(sel) / max(len(sel), 1):.0%}')
    sel = [h for p, h in conf_hits if p >= 0.85]; prec = sum(sel) / max(len(sel), 1)
    snap = prec >= 0.9 and len(sel) >= 15        # snap decisions only once earned on builds it never saw
    m = fit(X, y); pickle.dump({'model': m, 'cols': cols, 'snap': snap, 'prec85': prec, 'n85': len(sel), 'top1': hits / n},
                               open(os.path.join(HERE, 'teacher.pkl'), 'wb')); print(f'saved teacher.pkl (snap decisions {"ON" if snap else "off"})')

_LOADED = None
def guess(turn, cands_desc):
    """probabilities for each candidate from the saved model, or None. cands_desc: [{'desc':..., 'name':..., 'aimed':...}]"""
    global _LOADED
    p = os.path.join(HERE, 'teacher.pkl')
    if _LOADED is None or _LOADED[0] != os.path.getmtime(p): _LOADED = (os.path.getmtime(p), pickle.load(open(p, 'rb')))
    M = _LOADED[1]; F = [features(turn, c) for c in cands_desc]
    X = np.array([[f.get(k, 0.0) for k in M['cols']] for f in F]); return softmax(M['model'].predict(X)), M