import json, os, pickle, re, numpy as np, onnxruntime as ort
from tokenizers import Tokenizer
_TP = re.compile(r"(?u)\b\w\w+\b")
class Strat:
    """strat (fast kNN): tier-1 tf-idf, escalating to tier-2 int8 student + class-restricted kNN memory. CPU, 1 thread."""
    def __init__(s, path, threads=1):
        c = json.load(open(os.path.join(path, "config.json"))); s.c = c; s.labels = c["labels"]; C = len(s.labels)
        s.wv, s.wi, s.cv, s.ci, s.cT, s.b = pickle.load(open(os.path.join(path, "fast.pkl"), "rb")); s.off = len(s.wv)
        so = ort.SessionOptions(); so.intra_op_num_threads = threads; so.inter_op_num_threads = 1
        s.sess = ort.InferenceSession(os.path.join(path, "student_int8.onnx"), so, providers=["CPUExecutionProvider"])
        s.bt = Tokenizer.from_file(os.path.join(path, "tokenizer.json")); s.bt.enable_truncation(c["maxlen"]); s.bt.no_padding()
        k = np.load(os.path.join(path, "knn.npz")); E, y = k["E"], k["y"]
        o = np.argsort(y, kind="stable"); s.Ei = np.ascontiguousarray(E[o]); s.yi = y[o].astype(np.int64)   # sorted by class
        s.coff = np.concatenate([[0], np.cumsum(np.bincount(s.yi, minlength=C))]).astype(np.int64)
        s.M = int(c.get("knn_classes", 20)); s.nl = C
    def _tf(s, terms, vocab, idf, off):
        cnt = {}
        for t in terms:
            j = vocab.get(t)
            if j is not None: cnt[j] = cnt.get(j, 0) + 1
        if not cnt: return None
        idx = np.fromiter(cnt.keys(), dtype=np.int64, count=len(cnt)); tf = np.fromiter(cnt.values(), dtype=np.float32, count=len(cnt))
        v = (1.0 + np.log(tf)) * idf[idx]; n = np.sqrt((v * v).sum())
        return idx + off, (v / n if n > 0 else v)
    def feat(s, text):
        t = text.lower(); w = _TP.findall(t)
        a = s._tf(w + [w[i] + " " + w[i + 1] for i in range(len(w) - 1)], s.wv, s.wi, 0)
        cg = []
        for wd in t.split():
            wd = " " + wd + " "; L = len(wd)
            for n in range(2, 6):
                o = 0; cg.append(wd[o:o + n])
                while o + n < L:
                    o += 1; cg.append(wd[o:o + n])
                if o == 0: break
        b = s._tf(cg, s.cv, s.ci, s.off)
        ps = [x for x in (a, b) if x is not None]
        if not ps: return np.zeros(0, np.int64), np.zeros(0, np.float32)
        return np.concatenate([p[0] for p in ps]), np.concatenate([p[1] for p in ps]).astype(np.float32)
    def fast_p(s, text):
        idx, val = s.feat(text)
        z = val @ s.cT[idx] + s.b
        z = z - z.max(); p = np.exp(z); return p / p.sum()
    def _knn(s, e, p, p_tf):
        if s.M <= 0 or s.M >= s.nl:
            sim = s.Ei @ e; lab = s.yi
        else:
            cand = np.argpartition(-p, s.M)[:s.M]
            if p_tf is not None: cand = np.union1d(cand, np.argpartition(-p_tf, 4)[:4])
            co = s.coff; sims = [s.Ei[co[c]:co[c + 1]] @ e for c in cand]
            sim = np.concatenate(sims); lab = np.concatenate([s.yi[co[c]:co[c + 1]] for c in cand])
        top = np.argpartition(-sim, 10)[:10]; w = np.exp((sim[top] - sim[top].max()) / 0.05)
        pk = np.bincount(lab[top], weights=w, minlength=s.nl); return pk / pk.sum()
    def student_p(s, text, p_tf=None):
        ids = np.array([s.bt.encode(text).ids], dtype=np.int64)
        lg, e = s.sess.run(None, {"input_ids": ids, "attention_mask": np.ones_like(ids)})
        z = lg[0] - lg[0].max(); p = np.exp(z); p /= p.sum()
        lam = s.c["lam"]
        if lam > 0: p = (1 - lam) * p + lam * s._knn(e[0], p, p_tf)
        mu = s.c.get("mu", 0.0)
        if mu > 0 and p_tf is not None: p = (1 - mu) * p + mu * p_tf
        return p
    def predict_one(s, text):
        p = s.fast_p(text)
        if p.max() >= s.c["tau"]: return s.labels[int(p.argmax())], "fast"
        return s.labels[int(s.student_p(text, p).argmax())], "student"