import argparse
import json
import math
import os
import shutil
import stat
import subprocess
import tempfile
import time
from pathlib import Path

import numpy as np
import pandas as pd

ROOT = Path(__file__).parent
BETA, GATE, K_HITS = 0.7, 20, 30
EXACT = 'ACDEFGHIKLMNPQRSTVWY'
GROUPS = ['LVIM', 'C', 'A', 'G', 'ST', 'P', 'FYW', 'EDNQ', 'KR', 'H']
K, BASE = 6, 11
MAXDF, MAXPOS, BAND, TOP = 2000, 60000, 32, 200
KEYMUL = 2 * MAXPOS // BAND + 2
LETTERS = np.frombuffer(b'ACDEFGHIKLMNPQRSTVWYX', dtype=np.uint8)
COLS = ['q', 't', 'bits', 'pident', 'qcov', 'tcov']
NF = 29

LUT = np.full(256, 20, dtype=np.int64)
for i, c in enumerate(EXACT):
    LUT[ord(c)] = i
MAP = np.full(21, 10, dtype=np.int64)
for g, letters in enumerate(GROUPS):
    for c in letters:
        MAP[EXACT.index(c)] = g
LUT_NET = np.full(256, 21, dtype=np.int64)
for i, c in enumerate(EXACT):
    LUT_NET[ord(c)] = i + 1


def to_codes(seq):
    return LUT[np.frombuffer(seq.encode('ascii', 'ignore'), dtype=np.uint8)]


def hits_of(codes):
    codes = MAP[codes[:MAXPOS].astype(np.int64)]
    n = len(codes) - K + 1
    if n <= 0:
        return np.empty(0, dtype=np.int32), np.empty(0, dtype=np.int32)
    ids = np.zeros(n, dtype=np.int64)
    for j in range(K):
        ids = ids * BASE + codes[j:j + n]
    return ids.astype(np.int32), np.arange(n, dtype=np.int32)


def build_index(codes, lens):
    starts = np.cumsum(lens) - lens
    parts = [hits_of(codes[s:s + n]) for s, n in zip(starts, lens)]
    sizes = np.array([len(p[0]) for p in parts])
    ids_all = np.concatenate([p[0] for p in parts])
    pos_all = np.concatenate([p[1] for p in parts]).astype(np.uint16)
    del parts
    order = np.argsort(ids_all, kind='stable')
    kmer = ids_all[order]
    del ids_all
    pos = pos_all[order]
    del pos_all
    seq = np.repeat(np.arange(len(lens), dtype=np.int32), sizes)[order]
    return {'kmer': kmer, 'pos': pos, 'seq': seq}


def candidates(index, codes):
    empty = np.empty(0, dtype=np.int64)
    ids, qpos = hits_of(codes)
    if len(ids) == 0:
        return empty, empty
    lo = np.searchsorted(index['kmer'], ids, 'left')
    cnt = np.searchsorted(index['kmer'], ids, 'right') - lo
    keep = cnt <= MAXDF
    lo, cnt, qpos = lo[keep], cnt[keep], qpos[keep]
    total = cnt.sum()
    if total == 0:
        return empty, empty
    offs = np.cumsum(cnt) - cnt
    idx = np.arange(total) - np.repeat(offs, cnt) + np.repeat(lo, cnt)
    diag = index['pos'][idx].astype(np.int64) - np.repeat(qpos, cnt).astype(np.int64) + MAXPOS
    key = index['seq'][idx].astype(np.int64) * KEYMUL + diag // BAND
    uk, counts = np.unique(key, return_counts=True)
    nxt = np.minimum(np.searchsorted(uk, uk + 1), len(uk) - 1)
    score = counts + np.where(uk[nxt] == uk + 1, counts[nxt], 0)
    top = np.argsort(-score)[:TOP]
    cand = uk[top] // KEYMUL
    _, first = np.unique(cand, return_index=True)
    first.sort()
    return cand[first], score[top][first]


def vote(cand, score, lab_flat, lab_ptr, n_ec, k=5):
    v = np.zeros(n_ec, dtype=np.float32)
    for j, s in zip(cand[:k], score[:k]):
        labs = lab_flat[lab_ptr[j]:lab_ptr[j + 1]]
        v[labs] += s / len(labs)
    total = v.sum()
    return v / total if total > 0 else v


def red_votes(index, lab_flat, lab_ptr, seqs, n_ec):
    out = np.zeros((len(seqs), n_ec), dtype=np.float32)
    best = np.zeros(len(seqs))
    for r, seq in enumerate(seqs):
        cand, s = candidates(index, to_codes(seq))
        out[r] = vote(cand, s, lab_flat, lab_ptr, n_ec)
        best[r] = s[0] if len(s) else 0
    return out, best


def binary_path():
    flags = ''
    try:
        flags = Path('/proc/cpuinfo').read_text()
    except OSError:
        pass
    path = ROOT.resolve() / ('mmseqs_avx2' if ' avx2' in flags else 'mmseqs_sse41')
    try:
        os.chmod(path, os.stat(path).st_mode | stat.S_IEXEC)
    except OSError:
        pass
    return str(path)


def search(codes, lens, queries, threads=2, sens=4.0, max_seqs=50, timeout=480):
    work = tempfile.mkdtemp(dir='/tmp')
    try:
        exe = binary_path()
        starts = np.cumsum(lens) - lens
        with open(os.path.join(work, 'target.fa'), 'wb') as f:
            for i in range(len(lens)):
                seq = LETTERS[codes[starts[i]:starts[i] + lens[i]].astype(np.int64)].tobytes()
                f.write(b'>' + str(i).encode() + b'\n' + seq + b'\n')
        with open(os.path.join(work, 'query.fa'), 'wb') as f:
            for i, s in enumerate(queries):
                f.write(b'>' + str(i).encode() + b'\n' + s.encode('ascii', 'ignore') + b'\n')

        def run(*args):
            subprocess.run([exe, *args, '-v', '1'], check=True, cwd=work, timeout=timeout,
                           stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

        run('createdb', 'target.fa', 'tDB')
        run('createdb', 'query.fa', 'qDB')
        run('search', 'qDB', 'tDB', 'res', 'tmp', '-s', str(sens), '--max-seqs', str(max_seqs), '-e', '10',
            '--threads', str(threads), '--split-memory-limit', '6G')
        run('convertalis', 'qDB', 'tDB', 'res', 'res.m8', '--format-output', 'query,target,bits,pident,qcov,tcov')
        out = os.path.join(work, 'res.m8')
        if os.path.getsize(out) == 0:
            return pd.DataFrame(columns=COLS)
        table = pd.read_csv(out, sep='\t', names=COLS)
        return table.sort_values(['q', 'bits'], ascending=[True, False], kind='stable')
    finally:
        shutil.rmtree(work, ignore_errors=True)


def build_vocab_maps(ec_list):
    l3 = {}
    l2 = {}
    l3of = np.zeros(len(ec_list), dtype=np.int32)
    l2of = np.zeros(len(ec_list), dtype=np.int32)
    for i, e in enumerate(ec_list):
        p = e.split('.')
        l3of[i] = l3.setdefault('.'.join(p[:3]), len(l3))
        l2of[i] = l2.setdefault('.'.join(p[:2]), len(l2))
    return l3of, l2of


def query_features(t, bits, pid, qcov, tcov, labs, prior_log, l3of, l2of, qlen, tlens):
    prior_cnt = np.exp(prior_log)
    K = len(t)
    top = float(bits[0])
    rel = bits / top
    r4 = rel ** 4
    r10 = rel ** 10
    cand = {}
    tot4 = 0.0
    tot1 = 0.0
    l3s = {}
    l2s = {}
    for r in range(K):
        L = labs[t[r]]
        n = len(L)
        w4 = r4[r] / n
        w1 = rel[r] / n
        tot4 += r4[r]
        tot1 += rel[r]
        for e in L:
            d = cand.get(e)
            if d is None:
                d = [0.0] * 14
                d[5] = r
                d[6] = bits[r]
                d[7] = pid[r]
                d[8] = n
                d[9] = qcov[r] * tcov[r]
                d[10] = tlens[t[r]] / max(qlen, 1)
                cand[e] = d
            d[0] += w1
            d[12] += pid[r]
            d[13] += bits[r]
            d[1] += w4
            d[2] += r10[r] / n
            d[3] += 1
            if r < 5:
                d[4] += 1
            if r < 10:
                d[11] += 1
        for e in set(int(l3of[x]) for x in L):
            l3s[e] = l3s.get(e, 0.0) + r4[r] / len(set(int(l3of[x]) for x in L))
        for e in set(int(l2of[x]) for x in L):
            l2s[e] = l2s.get(e, 0.0) + r4[r] / len(set(int(l2of[x]) for x in L))
    ecs = list(cand.keys())
    F = np.zeros((len(ecs), NF), dtype=np.float32)
    top_pid = float(pid[0])
    top_labs = set(labs[t[0]])
    for i, e in enumerate(ecs):
        d = cand[e]
        F[i, 0] = d[0] / tot1
        F[i, 1] = d[1] / tot4
        F[i, 2] = d[2]
        F[i, 3] = d[3]
        F[i, 4] = d[4]
        F[i, 5] = d[5]
        F[i, 6] = d[6] / top
        F[i, 7] = d[7]
        F[i, 8] = d[8]
        F[i, 9] = d[9]
        F[i, 10] = d[10]
        F[i, 11] = d[11]
        F[i, 12] = prior_log[e]
        F[i, 13] = l3s[int(l3of[e])] / tot4
        F[i, 14] = l2s[int(l2of[e])] / tot4
        F[i, 15] = 1.0 if e in top_labs else 0.0
        F[i, 16] = np.log(top)
        F[i, 17] = top_pid
        F[i, 18] = np.log(max(qlen, 1))
        F[i, 19] = K
        F[i, 20] = d[1]
        F[i, 21] = len(top_labs)
        F[i, 22] = bits[min(1, K - 1)] / top
        F[i, 23] = d[6]
        F[i, 24] = d[3] / prior_cnt[e]
        F[i, 25] = d[12] / d[3]
        F[i, 26] = d[13] / bits.sum()
        F[i, 27] = np.log(d[3] / prior_cnt[e] + 1e-4)
        F[i, 28] = pid.mean()
    return ecs, F


def full_features(F):
    v = F[:, 1]
    v0 = F[:, 0]
    srt = np.sort(v)[::-1]
    ex = np.zeros((len(F), 7), dtype=np.float32)
    ex[:, 0] = len(F)
    ex[:, 1] = v - srt[0]
    ex[:, 2] = v - (srt[1] if len(srt) > 1 else 0)
    ex[:, 3] = (-v).argsort().argsort()
    ex[:, 4] = v0 - v0.max()
    pos = v[v > 0]
    ex[:, 5] = -(pos * np.log(pos)).sum()
    ex[:, 6] = F[:, 12] - F[:, 12].max()
    return np.hstack([F, ex])


class Ranker:
    def __init__(self, z):
        self.feat, self.thr = z['r_feat'], z['r_thr']
        self.left, self.right = z['r_left'], z['r_right']
        self.val, self.leaf = z['r_val'], z['r_leaf']
        self.base = float(z['r_base'][0])

    def score(self, X):
        rows = np.arange(len(X))
        raw = np.full(len(X), self.base)
        for i in range(len(self.feat)):
            node = np.zeros(len(X), dtype=np.int64)
            for _ in range(64):
                lf = self.leaf[i][node]
                if lf.all():
                    break
                f = self.feat[i][node]
                go_left = X[rows, f] <= self.thr[i][node]
                nxt = np.where(go_left, self.left[i][node], self.right[i][node])
                node = np.where(lf, node, nxt)
            raw += self.val[i][node]
        return raw


def rank_queries(hits, seqs, labs, prior_log, l3of, l2of, tlens, ranker, ec_list):
    out = {}
    q = hits.q.values
    if len(q) == 0:
        return out
    cuts = np.concatenate([[0], np.where(q[1:] != q[:-1])[0] + 1, [len(q)]])
    t_all, b_all = hits.t.values, hits.bits.values.astype(float)
    p_all, qc_all, tc_all = hits.pident.values, hits.qcov.values, hits.tcov.values
    feats, owners = [], []
    for a, b in zip(cuts[:-1], cuts[1:]):
        b = min(b, a + K_HITS)
        qi = int(q[a])
        ecs, F = query_features(t_all[a:b], b_all[a:b], p_all[a:b], qc_all[a:b], tc_all[a:b], labs, prior_log,
                                l3of, l2of, len(seqs[qi]), tlens)
        feats.append(full_features(F))
        owners.append((qi, ecs))
    scores = ranker.score(np.vstack(feats))
    pos = 0
    for qi, ecs in owners:
        s = scores[pos:pos + len(ecs)]
        pos += len(ecs)
        out[qi] = ec_list[ecs[int(np.argmax(s))]]
    return out


def network_scores(z, seqs):
    import torch
    import torch.nn.functional as F

    torch.set_grad_enabled(False)
    torch.set_num_threads(2)
    P = {k[2:]: torch.from_numpy(z[k]) for k in z.files if k.startswith('p_')}
    dilations = z['dilations'].tolist()
    alpha = float(z['alpha'][0])
    bank = F.normalize(torch.from_numpy(z['bank'].astype(np.float32)), dim=1)
    bank_flat = z['bank_flat'].astype(np.int64)
    bank_ptr = np.concatenate([[0], np.cumsum(z['bank_len'])])
    n_ec = len(z['ec_list'])

    def softmax(x):
        e = torch.exp(x - x.max(1, keepdim=True).values)
        return e / e.sum(1, keepdim=True)

    def encode(seq):
        codes = LUT_NET[np.frombuffer(seq.encode('ascii', 'ignore'), dtype=np.uint8)]
        if len(codes) > 800:
            start = (len(codes) - 800) // 2
            codes = codes[start:start + 800]
        return codes

    def forward(tokens):
        m = (tokens > 0).unsqueeze(-1).to(P['embed.w'].dtype)
        x = P['embed.w'][tokens] * m
        for i, d in enumerate(dilations):
            pre = f'b{i}.'
            w = P[pre + 'dw.w']
            xc = x.transpose(1, 2).contiguous()
            h = F.conv1d(xc, w, P[pre + 'dw.b'], padding=d * (w.shape[-1] // 2), dilation=d, groups=xc.shape[1])
            h = h.transpose(1, 2).contiguous() * m
            mu = h.mean(-1, keepdim=True)
            var = h.var(-1, unbiased=False, keepdim=True)
            h = P[pre + 'ln.g'] * ((h - mu) * (1.0 / torch.sqrt(var + 1e-6))) + P[pre + 'ln.b']
            h = h @ P[pre + 'pw1.w'] + P[pre + 'pw1.b']
            h = h * 0.5 * (1.0 + torch.erf(h / math.sqrt(2.0)))
            h = h @ P[pre + 'pw2.w'] + P[pre + 'pw2.b']
            x = (x + h) * m
        n = m.sum(1).clamp(min=1.0)
        mean = x.sum(1) / n
        mx, idx = (x - (1.0 - m) * 1e4).max(1)
        pooled = torch.cat([mx, mean], 1)
        mu = pooled.mean(-1, keepdim=True)
        var = pooled.var(-1, unbiased=False, keepdim=True)
        z1 = P['head.ln.g'] * ((pooled - mu) * (1.0 / torch.sqrt(var + 1e-6))) + P['head.ln.b']
        e = z1 @ P['head.proj.w'] + P['head.proj.b']
        a = e * 0.5 * (1.0 + torch.erf(e / math.sqrt(2.0)))
        return a @ P['head.out.w'] + P['head.out.b'], e

    order = np.argsort([len(s) for s in seqs], kind='stable')
    emb = np.zeros((len(seqs), 128), np.float32)
    probs = np.zeros((len(seqs), n_ec), np.float32)
    for s in range(0, len(seqs), 64):
        ids = order[s:s + 64]
        enc = [encode(seqs[i]) for i in ids]
        x = np.zeros((len(enc), max(len(e) for e in enc)), dtype=np.int64)
        for r, e in enumerate(enc):
            x[r, :len(e)] = e
        logits, e = forward(torch.from_numpy(x))
        emb[ids] = e.cpu().numpy()
        probs[ids] = softmax(logits).cpu().numpy()
    q = F.normalize(torch.from_numpy(emb), dim=1)
    votes = np.zeros((len(q), n_ec), np.float32)
    for s in range(0, len(q), 256):
        v, idx = (q[s:s + 256] @ bank.T).topk(5, dim=1)
        w = softmax(v / 0.05).numpy()
        idx = idx.numpy()
        for r in range(len(idx)):
            for j in range(5):
                labs = bank_flat[bank_ptr[idx[r, j]]:bank_ptr[idx[r, j] + 1]]
                votes[s + r, labs] += w[r, j] / len(labs)
    return alpha * votes + (1 - alpha) * probs


def predict_file(input_path, output_path):
    t0 = time.time()
    frame = pd.read_parquet(input_path)[['id', 'sequence']]
    if frame.id.isna().any() or frame.id.duplicated().any():
        raise ValueError('Input IDs must be present and unique.')
    frame = frame.drop_duplicates('sequence', keep='first').reset_index(drop=True)
    seqs = frame.sequence.tolist()

    z = np.load(ROOT / 'model.npz')
    ec_list = z['ec_list'].tolist()
    n_ec = len(ec_list)
    codes, lens = z['codes'], z['lens']
    lab_flat = z['lab_flat'].astype(np.int64)
    lab_ptr = np.concatenate([[0], np.cumsum(z['lab_len'])])

    chosen = {}
    try:
        labs = [lab_flat[lab_ptr[i]:lab_ptr[i + 1]] for i in range(len(lens))]
        prior_log = np.log(np.maximum(np.bincount(lab_flat, minlength=n_ec), 1).astype(float))
        l3of, l2of = build_vocab_maps(ec_list)
        ranker = Ranker(z)
        hits = search(codes, lens, seqs)
        print('search finished after', round(time.time() - t0), 's, hits:', len(hits), flush=True)
        chosen = rank_queries(hits, seqs, labs, prior_log, l3of, l2of, lens, ranker, ec_list)
    except Exception as error:
        print('search or ranking failed, using the backup method for every input:', error, flush=True)
        chosen = {}
    rest = np.array([i for i in range(len(seqs)) if i not in chosen], dtype=np.int64)
    print('ranked inputs:', len(chosen), ' inputs without a search hit:', len(rest), ' of', len(seqs), flush=True)

    preds = [chosen.get(i) for i in range(len(seqs))]
    if len(rest):
        index = build_index(codes, lens)
        red, best = red_votes(index, lab_flat, lab_ptr, [seqs[i] for i in rest], n_ec)
        weak = np.where(best < GATE)[0]
        if len(weak):
            try:
                net = network_scores(z, [seqs[rest[i]] for i in weak])
                red[weak] = BETA * red[weak] + (1 - BETA) * net
            except Exception as error:
                print('network step skipped:', error, flush=True)
        for r, i in enumerate(rest):
            preds[i] = ec_list[int(red[r].argmax())]

    with Path(output_path).open('w') as output:
        for key, ec in zip(frame.id, preds):
            output.write(json.dumps({'id': key, 'prediction': [ec]}) + '\n')
    print('total time', round(time.time() - t0), 's', flush=True)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', required=True)
    parser.add_argument('--output', required=True)
    args = parser.parse_args()
    predict_file(args.input, args.output)