"""Evidence for the v2 reranker: one feature vector per (query, candidate EC).

Hits are read from a stream, not a file: DIAMOND's alignments run to gigabytes and the task forbids
predict.py from writing any file over 16 MiB. Each query is reduced to its features as it arrives.

Four groups of evidence:

  V1    the homology and hierarchy features of v1.
  HOM   per-EC alignment geometry: both coverages, how an EC's best hit compares with the best hit
        that does not carry it, and what share of an EC's known members the query reached.
  PROF  residue-level evidence (see `profile_features`), new in v2.
  EMB   the hierarchical ESM-2 head: L4 and L3 probabilities, centroid similarity, nearest
        neighbours in the contrastive space.
"""
import math
from collections import defaultdict

import numpy as np

V1_NAMES = ["vote_frac", "vote_rank", "max_bits", "max_bits_rel", "n_hits", "max_pid", "mean_pid",
            "max_cov", "best_hit_is_ec", "best_hit_n_ecs", "q_best_bits", "q_best_pid", "q_n_hits",
            "q_n_cands", "q_len", "l1_mass", "l2_mass", "l3_mass", "ec_prior"]
HOM_NAMES = ["scov_max", "len_ratio", "bits_gap", "top3_rel", "first_rank", "member_frac",
             "hit_frac", "log_n_db", "qcov_union", "share_in_l3", "n_sib_l3"]
PROF_NAMES = ["vote_hk", "prof_llr", "prof_llr_l3", "prof_top", "sdp_match", "sdp_n", "prof_self",
              "prof_n"]
EMB_NAMES = ["emb_cls", "emb_l3", "emb_cen", "emb_knn", "emb_cls_rank", "emb_cen_rank",
             "emb_knn_rank", "has_emb", "emb_new"]
FEATURE_NAMES = V1_NAMES + HOM_NAMES + PROF_NAMES + EMB_NAMES
N_SCAN = len(V1_NAMES) + len(HOM_NAMES) + len(PROF_NAMES)

OUTFMT = ["qseqid", "sseqid", "pident", "bitscore", "length", "qlen", "slen", "qstart", "qend",
          "sstart", "send", "qseq_gapped", "sseq_gapped"]

AA = "ACDEFGHIKLMNPQRSTVWY"
LUT = np.full(256, 20, np.int8)
for _i, _a in enumerate(AA):
    LUT[ord(_a)] = _i
# Robinson & Robinson amino-acid background, the usual pseudocount prior.
BG = np.array([.0780, .0193, .0536, .0630, .0386, .0738, .0220, .0515, .0574, .0902,
               .0224, .0449, .0520, .0426, .0513, .0712, .0584, .0644, .0133, .0322])
BG = BG / BG.sum()


def read_hits(lines, k, keep_aln):
    """Group DIAMOND tabular lines by query; only the top `keep_aln` hits keep their alignment."""
    cur, rows = None, []
    for line in lines:
        p = line.rstrip("\n").split("\t")
        if len(p) < 9:  # a line cut short when the deadline stopped DIAMOND
            continue
        q = int(p[0])
        if q != cur:
            if rows:
                yield cur, rows
            cur, rows = q, []
        if len(rows) >= k:
            continue
        complete = len(p) > 12 and len(p[11]) == len(p[12])
        aln = (p[11], p[12]) if complete and len(rows) < keep_aln else None
        rows.append((int(p[1]), float(p[2]), float(p[3]), int(p[4]), int(p[5]), int(p[6]),
                     int(p[7]), int(p[8]), aln))
    if rows:
        yield cur, rows


def project(rows, qlen):
    """Alignments as a hits x query-columns matrix: -1 unaligned, 0-19 residue, 20 gap or other."""
    A = np.full((len(rows), qlen), -1, np.int8)
    for h, r in enumerate(rows):
        q = np.frombuffer(r[8][0].encode(), np.uint8)
        s = np.frombuffer(r[8][1].encode(), np.uint8)
        col = LUT[s[q != 45]]
        A[h, r[6] - 1:r[6] - 1 + len(col)] = col
    return A


def profile_features(rows, seq, cands, ecsets, beta=1.0, n_top=20):
    """Residue-level evidence for each candidate EC, from a query-anchored alignment.

    The fourth EC digit is usually decided by a few specificity-determining residues, invisible to
    any global similarity score. Every hit is projected on the query's columns, then for candidate
    EC e the hits are split into members and non-members, each giving a profile (Henikoff-weighted
    so redundant reference sequences do not dominate). Then we ask which profile the query fits:

      prof_llr     mean log-odds of the query's residues, members vs non-members
      prof_llr_l3  the same against same-L3 ECs only: the competitors that matter for L4
      prof_top     the same on the 20 most discriminative columns (Jensen-Shannon divergence)
      sdp_match    of columns conserved in e and different elsewhere, the share the query matches
      sdp_n        how many such columns exist
      prof_self    log-odds of the query's residues against e's profile vs background
      vote_hk      redundancy-corrected vote share

    Returns {ec: [vote_hk, llr, llr_l3, top, sdp_match, sdp_n, self, n_members]}.
    """
    A = project(rows, len(seq))
    H, L = A.shape
    hh, cc = np.nonzero(A >= 0)
    if not len(hh):
        return {}
    aa = A[hh, cc].astype(np.int64)
    cell = cc * 21 + aa
    cnt = np.bincount(cell, minlength=L * 21)
    kinds = (cnt.reshape(L, 21) > 0).sum(1)
    w = np.bincount(hh, 1.0 / (kinds[cc] * cnt[cell]), minlength=H)
    w = w / np.maximum(np.bincount(hh, minlength=H), 1)
    w = w * H / max(w.sum(), 1e-9)
    cw = w[hh]
    total = np.bincount(cell, cw, minlength=L * 21).reshape(L, 21)[:, :20]
    qi = LUT[np.frombuffer(seq.encode(), np.uint8)].astype(np.int64)
    known = np.where(qi < 20)[0]
    wsum = w.sum()

    def freq(c, n):
        return (c + beta * BG) / (n + beta)[:, None]

    out = {}
    member = {e: np.array([e in s for s in ecsets]) for e in cands}
    for e in cands:
        m = member[e]
        if not m.any():
            continue
        vote_hk = float(w[m].sum() / wsum)
        sel = m[hh]
        ce = np.bincount(cell[sel], cw[sel], minlength=L * 21).reshape(L, 21)[:, :20]
        co = np.maximum(total - ce, 0.0)
        ne, no = ce.sum(1), co.sum(1)
        pe = freq(ce, ne)
        cols = known[ne[known] > 0]
        f = [vote_hk, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, float(m.sum())]
        if len(cols):
            f[6] = float(np.mean(np.log(pe[cols, qi[cols]] / BG[qi[cols]])))
        both = known[(ne[known] > 0) & (no[known] > 0)]
        if len(both):
            po = freq(co, no)
            llr = np.log(pe[both, qi[both]] / po[both, qi[both]])
            f[1] = float(llr.mean())
            mix = 0.5 * (pe[both] + po[both])
            js = 0.5 * (pe[both] * np.log(pe[both] / mix)).sum(1) \
                + 0.5 * (po[both] * np.log(po[both] / mix)).sum(1)
            f[3] = float(llr[np.argsort(-js)[:n_top]].mean())
            sdp = (js > 0.25) & (pe[both].max(1) >= 0.6)
            if sdp.any():
                f[4] = float(np.mean(pe[both[sdp]].argmax(1) == qi[both[sdp]]))
                f[5] = float(sdp.sum())
        pre = e.rsplit(".", 1)[0]
        sib = np.array([any(x != e and x.rsplit(".", 1)[0] == pre for x in s) and not m[h]
                        for h, s in enumerate(ecsets)])
        if sib.any():
            sel = sib[hh]
            cs = np.bincount(cell[sel], cw[sel], minlength=L * 21).reshape(L, 21)[:, :20]
            ns = cs.sum(1)
            c2 = known[(ne[known] > 0) & (ns[known] > 0)]
            if len(c2):
                ps = freq(cs, ns)
                f[2] = float(np.log(pe[c2, qi[c2]] / ps[c2, qi[c2]]).mean())
        out[e] = f
    return out


def query_features(rows, seq, db_ecs, ec_count, ec_prior, cfg, profile=True):
    """All homology-side features for one query: (candidates, matrix, query-level info)."""
    rows = sorted(rows, key=lambda r: -r[2])
    best_bits = rows[0][2]
    best_pid = max(r[1] for r in rows)
    ecsets = [set(db_ecs[r[0]]) for r in rows]
    best_ecs = ecsets[0]
    qlen = rows[0][4]
    power, rel_cut, max_cands = cfg["power"], cfg["rel_cut"], cfg["max_cands"]

    agg = {}
    for h, (r, ecs) in enumerate(zip(rows, ecsets)):
        s, pid, bits, length, _, slen, qstart, qend, _ = r
        if bits < rel_cut * best_bits:
            continue
        for ec in ecs:
            a = agg.get(ec)
            if a is None:
                a = agg[ec] = {"sum": 0.0, "bits": [], "pids": [], "max_cov": 0.0, "max_scov": 0.0,
                               "first": h, "spans": [], "len_ratio": slen / max(qlen, 1)}
            a["sum"] += bits ** power / len(ecs)
            a["bits"].append(bits)
            a["pids"].append(pid)
            a["max_cov"] = max(a["max_cov"], length / max(qlen, 1))
            a["max_scov"] = max(a["max_scov"], length / max(slen, 1))
            a["spans"].append((qstart, qend))
    total = sum(a["sum"] for a in agg.values()) or 1.0
    ranked = sorted(agg, key=lambda e: -agg[e]["sum"])[:max_cands]

    mass = [defaultdict(float) for _ in range(3)]
    l3_ecs = defaultdict(list)
    for ec, a in agg.items():
        parts = ec.split(".")
        for lv in range(3):
            mass[lv][".".join(parts[:lv + 1])] += a["sum"] / total
        l3_ecs[".".join(parts[:3])].append(ec)

    prof = {}
    if profile and len(ranked) > 1:
        kept = [h for h, r in enumerate(rows) if r[8] is not None]
        prof = profile_features([rows[h] for h in kept], seq, ranked[:cfg["prof_cands"]],
                                [ecsets[h] for h in kept])

    q_level = [best_bits, best_pid, len(rows), len(agg), len(seq)]
    X = np.zeros((len(ranked), N_SCAN), np.float32)
    for j, ec in enumerate(ranked):
        a = agg[ec]
        pre = ec.split(".")
        l3 = ".".join(pre[:3])
        v1 = [a["sum"] / total, j, a["bits"][0], a["bits"][0] / best_bits, len(a["bits"]),
              max(a["pids"]), float(np.mean(a["pids"])), a["max_cov"], float(ec in best_ecs),
              len(best_ecs)] + q_level + [
                  mass[0][pre[0]], mass[1][".".join(pre[:2])], mass[2][l3], ec_prior.get(ec, 0.0)]
        other = next((r[2] for r, s in zip(rows, ecsets) if ec not in s), 0.0)
        n_db = ec_count.get(ec, 0)
        spans, cover, end = sorted(a["spans"]), 0, 0
        for st, en in spans:
            if en > end:
                cover += en - max(st - 1, end)
                end = en
        hom = [a["max_scov"], a["len_ratio"], (a["bits"][0] - other) / best_bits,
               float(np.mean(a["bits"][:3])) / best_bits, math.log1p(a["first"]),
               min(len(a["bits"]) / max(n_db, 1), 1.0), len(a["bits"]) / len(rows),
               math.log1p(n_db), cover / max(qlen, 1),
               a["sum"] / max(sum(agg[x]["sum"] for x in l3_ecs[l3]), 1e-9), len(l3_ecs[l3])]
        X[j] = v1 + hom + prof.get(ec, [0.0] * len(PROF_NAMES))
    return ranked, X, {"q_level": q_level, "best_ecs": len(best_ecs),
                       "mass": [dict(m) for m in mass], "max_cands": max_cands}


def assemble(scanned, n, emb, ec_prior, ec_count, seqs, extra_from_emb=5):
    """Merge homology features with head evidence into the final (q, ec, X) arrays.

    scanned[i] is query_features' result or None; emb maps a query to {ec: (cls, l3, cen, knn)}.
    A query missing from `emb` gets zeroed columns and has_emb = 0. ECs only the head proposes join
    as candidates with zeroed homology columns and emb_new = 1.
    """
    out_q, out_ec, out_X = [], [], []
    i_rank, i_prior, i_ndb = (V1_NAMES.index("vote_rank"), V1_NAMES.index("ec_prior"),
                              len(V1_NAMES) + HOM_NAMES.index("log_n_db"))
    i_qlevel = V1_NAMES.index("q_best_bits")
    i_best = V1_NAMES.index("best_hit_n_ecs")
    i_mass = V1_NAMES.index("l1_mass")
    for i in range(n):
        res = scanned[i]
        ev = emb.get(i, {})
        cands, X = (res[0], res[1]) if res else ([], np.zeros((0, N_SCAN), np.float32))
        known = set(cands)
        new = [e for e in sorted(ev, key=lambda e: -(ev[e][0] + ev[e][2] + ev[e][3]))
               if e not in known][:extra_from_emb]
        if not cands and not new:
            continue
        if new:
            Y = np.zeros((len(new), N_SCAN), np.float32)
            Y[:, i_rank] = 25
            for j, e in enumerate(new):
                Y[j, i_prior] = ec_prior.get(e, 0.0)
                Y[j, i_ndb] = math.log1p(ec_count.get(e, 0))
                parts = e.split(".")
                if res:
                    info = res[2]
                    Y[j, i_qlevel:i_qlevel + 5] = info["q_level"]
                    Y[j, i_best] = info["best_ecs"]
                    for lv in range(3):
                        Y[j, i_mass + lv] = info["mass"][lv].get(".".join(parts[:lv + 1]), 0.0)
                else:
                    Y[j, V1_NAMES.index("q_len")] = len(seqs[i])
            X = np.vstack([X, Y]) if len(X) else Y
        names = list(cands) + new
        E = np.zeros((len(names), len(EMB_NAMES)), np.float32)
        if ev:
            rank = [{e: r for r, e in enumerate(sorted(ev, key=lambda e: -ev[e][k]))}
                    for k in (0, 2, 3)]
            for j, e in enumerate(names):
                c = ev.get(e, (0.0, 0.0, 0.0, 0.0))
                E[j] = [c[0], c[1], c[2], c[3], rank[0].get(e, 99), rank[1].get(e, 99),
                        rank[2].get(e, 99), 1.0, float(e in new)]
        else:
            E[:, 4:7] = 99
        out_q.extend([i] * len(names))
        out_ec.extend(names)
        out_X.append(np.hstack([X, E]))
    F = np.vstack(out_X) if out_X else np.zeros((0, len(FEATURE_NAMES)), np.float32)
    return np.array(out_q, np.int64), np.array(out_ec), F