"""Shared by training (notebook) and inference (predict.py): candidate features + a numpy tree-ensemble evaluator."""
import json, math
import numpy as np

FEATS = ['mm_share', 'mm_maxb_ratio', 'mm_maxb', 'mm_cnt', 'mm_rank', 'mm_maxpid', 'mm_maxqcov', 'mm_maxtcov', 'topbits', 'top_pid', 'qlen', 'nhits',
         'h_prob', 'h_rank', 'h_max', 'mm_l3share', 'h_l3mass', 'logprior', 'ncand_mm',
         'in_top1', 'top3_cnt', 'share_gap', 'share_rank', 'h_rel', 'h_gap', 'top1_nlab', 'sup_nlab', 'tb_per_len', 'mb_per_len', 'q_margin', 'q_entropy', 'q_nclose', 'bits_gap', 'pid_gap', 'cnt_frac', 'maxb_rank', 'k_share', 'k_max', 'k_cnt', 'k_top1',
         'n2_share', 'n2_gap', 'n2_rank', 'n_pur', 'n_top1', 'is_unseen', 'fam_nunseen', 'fam_nseen', 'fam_rank',
         'hs_prob', 'hs_rank_fam', 'hs_gap_fam', 'ps_sim', 'ps_rank_fam', 'ps_gap_fam', 'ps_gap_all',
         'mo_prob', 'mo_rank', 'mo_fam', 'mo_fam_rank', 'mo_gap_fam', 'mo_l3mass']
KNN_FEATS = ['k_share', 'k_max', 'k_cnt', 'k_top1']
CTX_FEATS = ['q_margin', 'q_entropy', 'q_nclose', 'bits_gap', 'pid_gap', 'cnt_frac', 'maxb_rank']
UNSEEN_FEATS = ['is_unseen', 'fam_nunseen', 'fam_nseen', 'fam_rank']     # ECs of the official vocabulary that have no training sequence
NB_FEATS = ['n2_share', 'n2_gap', 'n2_rank', 'n_pur', 'n_top1']      # from the training-set neighbourhood graph (precomputed train-vs-train MMseqs2)
HS_FEATS = ['hs_prob', 'hs_rank_fam', 'hs_gap_fam']                        # within-L3 head: P(4th digit | family)
PS_FEATS = ['ps_sim', 'ps_rank_fam', 'ps_gap_fam', 'ps_gap_all']           # supervised-contrastive EC prototypes
HIER_FEATS = HS_FEATS + PS_FEATS
MOTIF_FEATS = ['mo_prob', 'mo_rank', 'mo_fam', 'mo_fam_rank', 'mo_gap_fam', 'mo_l3mass']   # full-sequence k-mer (3-6) motif classifier; available for EVERY query
HEAD_FEATS = ['h_prob', 'h_rank', 'h_max', 'h_l3mass', 'h_rel', 'h_gap'] + KNN_FEATS + HIER_FEATS   # all of these need the ESM embedding
COLS = ['q', 'e'] + FEATS
NAN = float('nan')


class Ctx:
    """labels: list (per target sequence) of tuples of EC indices; ecs: EC strings; prior: per-EC counts."""
    def __init__(self, labels, ecs, prior):
        self.labels, self.ecs = labels, list(ecs)
        self.prior = np.asarray(prior, float); self.logprior = np.log1p(self.prior)
        ids = {}
        self.l3_id = np.array([ids.setdefault('.'.join(e.split('.')[:3]), len(ids)) for e in self.ecs]); self.nl3 = len(ids)
        self.top_ec = self.ecs[int(self.prior.argmax())]
        self.nb_e = self.nb_s = None
        # ECs with no database sequence ("unseen"), grouped by L3 family; ordered by their 4th field (numeric where possible)
        def key4(e):
            f = e.split('.')[-1]
            return (0, int(f), f) if f.isdigit() else (1, 0, f)
        un = [i for i in range(len(self.ecs)) if self.prior[i] <= 0]
        self.unseen_by_l3 = {}
        for i in sorted(un, key=lambda i: key4(self.ecs[i])): self.unseen_by_l3.setdefault(int(self.l3_id[i]), []).append(i)
        self.fam_members = [np.flatnonzero(self.l3_id == f) for f in range(self.nl3)]
        self.nunseen_l3 = np.zeros(self.nl3); self.nseen_l3 = np.zeros(self.nl3)
        for i in range(len(self.ecs)):
            if self.prior[i] <= 0: self.nunseen_l3[self.l3_id[i]] += 1
            else: self.nseen_l3[self.l3_id[i]] += 1
        self.use_unseen = False
        self.use_hier_cands = False
        self.use_motif_cands = False


    def set_nb(self, E, S):
        """E[j], S[j]: the top EC indices among database sequence j's own neighbours and their (bits^2-weighted) shares; S==0 marks empty slots."""
        self.nb_e = None if E is None else np.asarray(E); self.nb_s = None if S is None else np.asarray(S, np.float64)


def query_rows(qi, hit, qlen, h, ctx, TOPM=50, TOPH=10, knn=None, TOPK=5, TOPN=20, TOPN2C=3, h2=None, ps=None, mot=None, TOPMO=3):
    """hit: None or (t, bits, pident, qcov, tcov) arrays sorted by bits desc; h: None or head probabilities over all ECs."""
    if hit is not None:
        t, b, pid, qc, tc = (np.asarray(a)[:TOPM] for a in hit); tb = float(b[0]); tpid = float(pid[0])
    else:
        t = b = pid = qc = tc = np.empty(0); tb = 0.0; tpid = 0.0
    votes, maxb, cnt, rank, mpid, mqc, mtc = {}, {}, {}, {}, {}, {}, {}
    nlab = {}
    for r, (j, bj, pj, qj, tj) in enumerate(zip(t, b, pid, qc, tc)):
        l = ctx.labels[int(j)]
        for e in l:
            votes[e] = votes.get(e, 0.0) + float(bj) ** 2 / len(l); maxb[e] = max(maxb.get(e, 0.0), float(bj)); cnt[e] = cnt.get(e, 0) + 1; rank.setdefault(e, r); nlab[e] = nlab.get(e, 0) + len(l)
            mpid[e] = max(mpid.get(e, 0.0), float(pj)); mqc[e] = max(mqc.get(e, 0.0), float(qj)); mtc[e] = max(mtc.get(e, 0.0), float(tj))
    tot = sum(votes.values()) or 1.0
    l3v = {}
    for e, v in votes.items(): l3v[ctx.l3_id[e]] = l3v.get(ctx.l3_id[e], 0.0) + v / tot
    if h is not None:
        l3m = np.bincount(ctx.l3_id, weights=h, minlength=ctx.nl3); hmax = float(h.max())
        cands = set(votes) | set(np.argpartition(-h, TOPH)[:TOPH].tolist())
    else:
        l3m = None; hmax = NAN; cands = set(votes)
    if knn is not None:    # nearest training sequences in ESM-embedding space: (positions in ctx.labels, cosine similarities)
        ki, ks = knn; ks = np.asarray(ks, float); kw = np.exp((ks - ks.max()) * 20.0); kv, km, kc = {}, {}, {}
        for j, sj, wj in zip(ki, ks, kw):
            l = ctx.labels[int(j)]
            for e in l: kv[e] = kv.get(e, 0.0) + wj / len(l); km[e] = max(km.get(e, -1.0), float(sj)); kc[e] = kc.get(e, 0) + 1
        ktot = sum(kv.values()) or 1.0; ktop = float(ks.max()); cands |= set(sorted(kv, key=lambda e: -kv[e])[:TOPK])
    fam_rank = {}
    if ctx.use_unseen:      # one representative unseen EC per likely L3 family (top-2 by homology share, top-1 by head mass)
        fams = sorted(l3v, key=lambda f: -l3v[f])[:2]
        if l3m is not None: fams += [int(np.argmax(l3m))]
        for r_, f in enumerate(fams):
            fam_rank.setdefault(int(f), r_)
            u = ctx.unseen_by_l3.get(int(f))
            if u: cands.add(u[0])
    use_h2, use_ps = h2 is not None, ps is not None
    if ctx.use_hier_cands and (use_h2 or use_ps):     # best sibling (within-family head / prototype) inside the two most likely families
        fams2 = sorted(l3v, key=lambda f: -l3v[f])[:2]
        if l3m is not None: fams2 += [int(np.argmax(l3m))]
        for f in set(int(x) for x in fams2):
            mem = ctx.fam_members[f]
            if use_h2: cands.add(int(mem[int(np.argmax(h2[mem]))]))
            if use_ps:
                v = ps[mem]
                if np.isfinite(v).any(): cands.add(int(mem[int(np.nanargmax(v))]))
    use_mo = mot is not None
    if use_mo:
        mfam = np.bincount(ctx.l3_id, weights=mot, minlength=ctx.nl3)
        if ctx.use_motif_cands: cands |= set(np.argpartition(-mot, TOPMO)[:TOPMO].tolist())
    use_nb = ctx.nb_e is not None
    n2, pur_v, pur_w, nt1 = {}, {}, {}, None
    if use_nb and len(t):
        wsum = 0.0
        for r, (j, bj) in enumerate(zip(t[:TOPN], b[:TOPN])):
            j = int(j); w = float(bj) ** 2; sj = ctx.nb_s[j]
            if sj[0] <= 0:
                continue                                   # this database sequence has no neighbours of its own
            nd = dict(zip(ctx.nb_e[j].tolist(), sj.tolist())); wsum += w
            if r == 0: nt1 = nd
            for e, v in nd.items():
                if v > 0: n2[e] = n2.get(e, 0.0) + w * v
            for e in ctx.labels[j]:
                pur_v[e] = pur_v.get(e, 0.0) + w * nd.get(e, 0.0); pur_w[e] = pur_w.get(e, 0.0) + w
        if wsum > 0: n2 = {e: v / wsum for e, v in n2.items()}
        cands |= set(sorted(n2, key=lambda e: -n2[e])[:TOPN2C])
    top1 = set(ctx.labels[int(t[0])]) if len(t) else set()
    top3 = {}
    for j in t[:3]:
        for e in ctx.labels[int(j)]: top3[e] = top3.get(e, 0) + 1
    shares = {e: votes.get(e, 0.0) / tot for e in cands}; smax = max(shares.values()) if shares else 0.0
    srank = {e: r for r, e in enumerate(sorted(shares, key=lambda x: -shares[x]))}
    mbs = {e: maxb.get(e, 0.0) for e in cands}; mbmax = max(mbs.values()) if mbs else 0.0
    mps = {e: mpid.get(e, 0.0) for e in cands}; mpmax = max(mps.values()) if mps else 0.0
    brank = {e: r for r, e in enumerate(sorted(mbs, key=lambda x: -mbs[x]))}
    sv = sorted(shares.values(), reverse=True)
    qmargin = (sv[0] - (sv[1] if len(sv) > 1 else 0.0)) if sv else 0.0
    pp = np.array([v for v in sv if v > 0]); qent = float(-(pp / pp.sum() * np.log(pp / pp.sum())).sum()) if len(pp) else 0.0
    nclose = sum(1 for v in sv if sv[0] > 0 and v >= 0.5 * sv[0]); nt = max(len(t), 1)
    n2c = {e: n2.get(e, 0.0) for e in cands}; n2max = max(n2c.values()) if n2c else 0.0
    n2rank = {e: r for r, e in enumerate(sorted(n2c, key=lambda x: -n2c[x]))}
    mor, mofr, mofg = {}, {}, {}
    if use_mo:
        for e in cands: mor[e] = min(int((mot > mot[e]).sum()), 99)
        for f in set(int(ctx.l3_id[e]) for e in cands):
            mem = ctx.fam_members[f]; v = mot[mem]; mx = float(v.max())
            for e in cands:
                if int(ctx.l3_id[e]) == f: mofr[e] = int((v > mot[e]).sum()); mofg[e] = (float(mot[e]) - mx) / max(float(mfam[f]), 1e-12)
    hsr, hsg, psr, psg, psa = {}, {}, {}, {}, {}
    if use_h2 or use_ps:
        for f in set(int(ctx.l3_id[e]) for e in cands):
            mem = ctx.fam_members[f]
            if use_h2:
                v = h2[mem]; mx = float(v.max()); order = {int(e): r for r, e in enumerate(mem[np.argsort(-v, kind='stable')])}
                for e in cands:
                    if int(ctx.l3_id[e]) == f: hsr[e] = order[int(e)]; hsg[e] = float(h2[e]) - mx
            if use_ps:
                v = ps[mem]; fin = np.isfinite(v)
                if fin.any():
                    mx = float(np.nanmax(v)); vv = np.where(fin, v, -np.inf); order = {int(e): r for r, e in enumerate(mem[np.argsort(-vv, kind='stable')])}
                    for e in cands:
                        if int(ctx.l3_id[e]) == f and np.isfinite(ps[e]): psr[e] = order[int(e)]; psg[e] = float(ps[e]) - mx
        if use_ps:
            pc = [float(ps[e]) for e in cands if np.isfinite(ps[e])]; pcmax = max(pc) if pc else NAN
    rows = []
    for e in cands:
        rows.append((qi, e, votes.get(e, 0.0) / tot, maxb.get(e, 0.0) / tb if tb > 0 else 0.0, maxb.get(e, 0.0), cnt.get(e, 0), rank.get(e, 99),
                     mpid.get(e, 0.0), mqc.get(e, 0.0), mtc.get(e, 0.0), tb, tpid, qlen, len(t),
                     float(h[e]) if h is not None else NAN, int((h > h[e]).sum()) if h is not None else NAN, hmax,
                     l3v.get(ctx.l3_id[e], 0.0), float(l3m[ctx.l3_id[e]]) if h is not None else NAN, float(ctx.logprior[e]), len(votes),
                     int(e in top1), top3.get(e, 0), shares[e] - smax, srank[e],
                     float(h[e]) / hmax if h is not None and hmax > 0 else NAN, float(h[e]) - hmax if h is not None else NAN,
                     len(ctx.labels[int(t[0])]) if len(t) else 0, nlab.get(e, 0) / cnt[e] if cnt.get(e, 0) else 0.0,
                     tb / max(qlen, 1), maxb.get(e, 0.0) / max(qlen, 1),
                     qmargin, qent, nclose, mbs[e] - mbmax, mps[e] - mpmax, cnt.get(e, 0) / nt, brank[e],
                     kv.get(e, 0.0) / ktot if knn is not None else NAN, km.get(e, -1.0) if knn is not None else NAN, kc.get(e, 0) if knn is not None else NAN, ktop if knn is not None else NAN,
                     n2c[e] if use_nb else NAN, n2c[e] - n2max if use_nb else NAN, n2rank[e] if use_nb else NAN,
                     pur_v[e] / pur_w[e] if use_nb and pur_w.get(e, 0) > 0 else NAN,
                     (nt1.get(e, 0.0) if nt1 is not None else NAN) if use_nb else NAN,
                     int(ctx.prior[e] <= 0), float(ctx.nunseen_l3[ctx.l3_id[e]]), float(ctx.nseen_l3[ctx.l3_id[e]]), fam_rank.get(int(ctx.l3_id[e]), 9),
                     float(h2[e]) if use_h2 else NAN, hsr.get(e, NAN) if use_h2 else NAN, hsg.get(e, NAN) if use_h2 else NAN,
                     float(ps[e]) if use_ps else NAN, psr.get(e, NAN) if use_ps else NAN, psg.get(e, NAN) if use_ps else NAN,
                     (float(ps[e]) - pcmax) if use_ps and np.isfinite(ps[e]) else NAN,
                     float(mot[e]) if use_mo else NAN, mor.get(e, NAN) if use_mo else NAN,
                     float(mot[e]) / max(float(mfam[ctx.l3_id[e]]), 1e-12) if use_mo else NAN, mofr.get(e, NAN) if use_mo else NAN,
                     mofg.get(e, NAN) if use_mo else NAN, float(mfam[ctx.l3_id[e]]) if use_mo else NAN))
    return rows


# ---------- LightGBM model -> arrays -> numpy prediction (no lightgbm needed at inference) ----------
def forest_from_lgb(booster):
    out = []
    for ti in booster.dump_model()['tree_info']:
        feat, thr, left, right, dl, mnan, val = [], [], [], [], [], [], []
        def add(node):
            k = len(feat); feat.append(-1); thr.append(0.0); left.append(-1); right.append(-1); dl.append(False); mnan.append(False); val.append(0.0)
            if 'leaf_value' in node: val[k] = float(node['leaf_value']); return k
            assert node['decision_type'] == '<=' and node['missing_type'] in ('None', 'NaN'), node['missing_type']
            feat[k] = int(node['split_feature']); thr[k] = float(node['threshold']); dl[k] = bool(node['default_left']); mnan[k] = node['missing_type'] == 'NaN'
            left[k] = add(node['left_child']); right[k] = add(node['right_child']); return k
        add(ti['tree_structure'])
        out.append(dict(feat=feat, thr=thr, left=left, right=right, dl=dl, mnan=mnan, val=val))
    return out


class Forest:
    def __init__(self, trees):
        self.trees = [{k: np.asarray(v, dtype={'feat': np.int64, 'left': np.int64, 'right': np.int64, 'dl': bool, 'mnan': bool}.get(k, np.float64)) for k, v in t.items()} for t in trees]
    def predict(self, X):
        X = np.asarray(X, np.float64); n = len(X); out = np.zeros(n)
        for tr in self.trees:
            node = np.zeros(n, np.int64)
            while True:
                idx = np.where(tr['feat'][node] >= 0)[0]
                if len(idx) == 0: break
                nd = node[idx]; x = X[idx, tr['feat'][nd]]; isn = np.isnan(x)
                left = np.where(isn, np.where(tr['mnan'][nd], tr['dl'][nd], 0.0 <= tr['thr'][nd]), x <= tr['thr'][nd])
                node[idx] = np.where(left, tr['left'][nd], tr['right'][nd])
            out += tr['val'][node]
        return out


# ---------- small helpers for the head ----------
def gelu(x):
    from scipy.special import erf
    return 0.5 * x * (1.0 + erf(x / math.sqrt(2.0)))

def softmax(z):
    z = z - z.max(1, keepdims=True); e = np.exp(z); return e / e.sum(1, keepdims=True)


def family_softmax(z, l3_id, nl3):
    """z: logits over all ECs, 1-D or (n, n_ECs). -> P(EC | its own L3 family) for every EC."""
    z = np.asarray(z, np.float64); one = z.ndim == 1; Z = z[None] if one else z
    ez = np.exp(Z - Z.max(1, keepdims=True))
    S = np.zeros((len(Z), nl3)); np.add.at(S.T, l3_id, ez.T)
    out = ez / np.maximum(S[:, l3_id], 1e-300)
    return out[0] if one else out


# ---------- full-sequence k-mer motif model (hashed k-mers -> mean embedding -> linear over ECs) ----------
_KLUT = np.full(256, 20, np.int64)
for _i, _a in enumerate('ACDEFGHIKLMNPQRSTVWY'): _KLUT[ord(_a)] = _i

def kmer_ids(seq, ks=(3, 4, 5, 6), nb=2 ** 20, maxlen=1000):
    """Hashed ids of every k-mer (k in ks) of the first maxlen residues. Identical in training and inference."""
    c = _KLUT[np.frombuffer(seq[:maxlen].encode('ascii', 'ignore'), np.uint8)]
    out = []
    for k in ks:
        n = len(c) - k + 1
        if n <= 0: continue
        code = np.zeros(n, np.int64)
        for j in range(k): code = code * 21 + c[j:j + n]
        out.append((code * 2654435761 + k * 97531) % nb)
    return np.concatenate(out).astype(np.int64) if out else np.zeros(1, np.int64)

def motif_probs(seqs, mw, ks=(3, 4, 5, 6), batch=1000):
    """mw: dict(table (nb, D) float16, W (nEC, D), b (nEC)). -> (n, nEC) softmax probabilities."""
    T = mw['table']; W = mw['W'].astype(np.float32); b = mw['b'].astype(np.float32); nb = T.shape[0]
    out = np.zeros((len(seqs), len(b)), np.float32)
    for st in range(0, len(seqs), batch):
        X = np.stack([T[kmer_ids(s, ks, nb)].astype(np.float32).mean(0) for s in seqs[st:st + batch]])
        out[st:st + batch] = softmax(X @ W.T + b)
    return out