"""Offline CPU inference: MMseqs2 search -> (time-budgeted) ESM-2 head on least-confident queries -> LightGBM-style reranker (numpy)."""
import argparse, gzip, json, os, pickle, shutil, subprocess, sys, tempfile, time
from pathlib import Path
import numpy as np
import pandas as pd

T0 = time.time()
ROOT = Path(__file__).resolve().parent
WDIR = Path(os.environ.get('EC_WEIGHTS', str(ROOT / 'weights')))   # override only used by our own end-to-end check
sys.path.insert(0, str(ROOT))
import ecpipe


def log(*a):
    print(f'[{time.time() - T0:6.1f}s]', *a, file=sys.stderr, flush=True)


# ---------------------------------------------------------------- diagnostics
def diagnostics():
    try:
        import platform
        cpu = open('/proc/cpuinfo').read() if os.path.exists('/proc/cpuinfo') else ''
        flags = [f for f in ('avx2', 'sse4_1', 'avx512f') if (' ' + f + ' ') in cpu or (' ' + f + '\n') in cpu]
        log('python', sys.version.split()[0], '|', platform.platform(), '| cpus', os.cpu_count(), '| uid', os.getuid(), '| cpu flags', flags)
        for line in open('/proc/mounts').read().splitlines():
            p = line.split()
            if len(p) > 3 and p[1] in ('/', '/tmp', '/dev/shm', '/app', '/var/tmp', '/home'): log('mount', p[1], p[2], p[3][:80])
        for d in ('/tmp', '/dev/shm', '/var/tmp', str(Path.home()), str(ROOT), os.getcwd()):
            try: log('dir', d, 'writable' if os.access(d, os.W_OK) else 'NOT writable', 'free GB', round(shutil.disk_usage(d).free / 1e9, 1))
            except Exception as e: log('dir', d, 'err', repr(e)[:80])
        for f in ('/sys/fs/cgroup/memory.max', '/sys/fs/cgroup/memory/memory.limit_in_bytes'):
            if os.path.exists(f): log('cgroup mem', open(f).read().strip()); break
        mi = open('/proc/meminfo').read().splitlines(); log('meminfo', [x for x in mi if x.startswith(('MemTotal', 'MemAvailable'))])
        for p in sorted(WDIR.rglob('*')):
            if p.is_file() and p.stat().st_size > 1e6: log('weights', str(p.relative_to(WDIR)), round(p.stat().st_size / 1e6), 'MB', oct(p.stat().st_mode & 0o777))
    except Exception as e:
        log('diagnostics failed', repr(e)[:200])


# ---------------------------------------------------------------- MMseqs2
# The grader refuses any file over 16 MiB (that is what broke the first version: copying the 17 MB binary raised EFBIG).
# So: run the binary in place, and search in small chunks so every file MMseqs2 writes stays far below the limit.
MM_CHUNK = int(os.environ.get('EC_MM_CHUNK', 400))   # queries per MMseqs2 call
MM_DEADLINE_S = 260                                  # no new first-pass MMseqs2 chunk starts after this many seconds
MM2_DEADLINE_S = 330                                 # no new second-pass (sensitive) chunk starts after this many seconds
EMBED_DEADLINE_S = 460                               # PLM embedding stops here
KNN_DEADLINE_S = 530                                 # embedding-space nearest-neighbour search stops here
MOTIF_DEADLINE_S = 520                               # full-sequence motif scoring stops here (features become missing)


class Runner:
    """Finds a way to run the bundled static mmseqs binary: in place (chmod if needed) first, copies only as a last resort."""
    def __init__(self):
        self.cmd0, self.fds, self.how = None, (), None
        try: cpu = open('/proc/cpuinfo').read()
        except Exception: cpu = ''
        names = ['mmseqs-avx2', 'mmseqs-sse41'] if ' avx2' in cpu else ['mmseqs-sse41', 'mmseqs-avx2']
        for name in names:
            src = WDIR / 'bin' / name
            if not src.exists(): continue
            if self._try(lambda: self._inplace(src), f'{name} in place'): return
            for d in ('/tmp', '/dev/shm', '/var/tmp', str(Path.home()), str(ROOT), os.getcwd()):
                if self._try(lambda d=d: self._copy(src, d), f'{name} copied to {d}'): return
        log('NO RUNNABLE mmseqs binary found')

    def _inplace(self, src):
        if not os.access(src, os.X_OK): os.chmod(src, 0o755)
        return [str(src)], ()
    def _copy(self, src, d):
        dst = Path(tempfile.mkdtemp(dir=d)) / 'mmseqs'; shutil.copy(src, dst); os.chmod(dst, 0o755); return [str(dst)], ()

    def _try(self, make, label):
        try:
            cmd0, fds = make()
            r = subprocess.run(cmd0 + ['version'], capture_output=True, text=True, timeout=60, pass_fds=fds)
            if r.returncode == 0:
                self.cmd0, self.fds, self.how = cmd0, fds, label; log('mmseqs runnable:', label, r.stdout.strip()[:30]); return True
            log('mmseqs not runnable:', label, 'rc', r.returncode, (r.stdout + r.stderr)[-200:].strip())
        except Exception as e:
            log('mmseqs not runnable:', label, repr(e)[:200])
        return False

    def run(self, args, timeout):
        r = subprocess.run(self.cmd0 + args, capture_output=True, text=True, timeout=max(20, timeout), pass_fds=self.fds)
        if r.returncode: raise RuntimeError(f'rc {r.returncode}: ' + (r.stdout + r.stderr)[-700:])
        return r


def pick_workdir(need_gb=0.3):
    for d in ('/tmp', '/dev/shm', '/var/tmp', str(Path.home()), str(ROOT), os.getcwd()):
        try:
            if os.access(d, os.W_OK) and shutil.disk_usage(d).free > need_gb * 1e9: return tempfile.mkdtemp(dir=d)
        except Exception: pass
    return tempfile.mkdtemp()


RUNNER = None


def parse_m8(path):
    if not os.path.exists(path) or os.path.getsize(path) == 0: return {}
    m8 = pd.read_csv(path, sep='\t', header=None, names=['q', 't', 'bits', 'pid', 'qc', 'tc'])
    if os.environ.get('EC_HIT_CAP'): m8 = m8[m8.pid < float(os.environ['EC_HIT_CAP'])]      # only for the notebook's end-to-end check on an identity-capped holdout
    m8 = m8.sort_values(['q', 'bits'], ascending=[True, False], kind='stable')
    q = m8.q.values; cut = np.flatnonzero(np.diff(q)) + 1; starts = np.r_[0, cut]; ends = np.r_[cut, len(q)]
    cols = [m8[c].values for c in ('t', 'bits', 'pid', 'qc', 'tc')]
    return {int(q[s]): tuple(c[s:e] for c in cols) for s, e in zip(starts, ends)}


def mm_hits(seqs, meta, deadline, idxs=None, s=None, max_seqs=None, chunk=None, tag='mmseqs', budget=None):
    """-> {query index: (target idx, bits, pident, qcov, tcov) sorted by bits desc}. Partial results are kept if later chunks fail (reasons are logged)."""
    hits = {}
    idxs = list(range(len(seqs))) if idxs is None else list(idxs)
    s = meta['mm_s'] if s is None else s; max_seqs = meta['mm_max_seqs'] if max_seqs is None else max_seqs; chunk = chunk or MM_CHUNK
    try:
        rn = RUNNER if RUNNER is not None else Runner()
        globals()['RUNNER'] = rn
        if rn.cmd0 is None: return {}
        db = str(WDIR / 'mm' / 'targetDB'); fails = 0; t_all = time.time(); degraded = False
        for st in range(0, len(idxs), chunk):
            if time.time() > deadline: log(f'{tag} deadline: stopping at {st}/{len(idxs)}'); break
            if budget and st >= 0.25 * len(idxs) and not degraded and (time.time() - t_all) / (st / len(idxs)) > budget:
                degraded = True; s = min(float(s), 2.5); max_seqs = min(int(max_seqs), 50)      # the grader is slow today: finish the remaining queries with a cheaper search
                log(f'{tag}: projected {(time.time() - t_all) / (st / len(idxs)):.0f}s > budget {budget}s, switching to s={s} max-seqs {max_seqs} for the rest')
            work = pick_workdir(); t = time.time(); rem = lambda: deadline - time.time() + 60
            try:
                with open(work + '/q.fa', 'w') as f:
                    for i in idxs[st:st + chunk]: f.write('>%d\n%s\n' % (i, seqs[i]))
                rn.run(['createdb', work + '/q.fa', work + '/qdb', '-v', '1'], rem())
                rn.run(['search', work + '/qdb', db, work + '/res', work + '/tmp', '-s', str(s), '--max-seqs', str(max_seqs),
                        '--threads', '2', '--split-memory-limit', '10G', '-v', '1'], rem())
                rn.run(['convertalis', work + '/qdb', db, work + '/res', work + '/res.m8', '--format-output', 'query,target,bits,pident,qcov,tcov', '-v', '1'], rem())
                hits.update(parse_m8(work + '/res.m8')); fails = 0
            except Exception as e:
                fails += 1; log(f'MMSEQS chunk at query {st} FAILED ({fails}):', repr(e)[:800])
                if fails >= 3: log('giving up on MMseqs2'); break
            finally:
                shutil.rmtree(work, ignore_errors=True)
            if (st // chunk) % 4 == 0: log(f'{tag} chunk {st // chunk} took {time.time() - t:.1f}s; hits so far {len(hits)}')
        log(f'{tag} total {time.time() - t_all:.0f}s')
    except Exception as e:
        log('MMSEQS FAILED, continuing with what we have:', repr(e)[:900])
    return hits


# ---------------------------------------------------------------- ESM-2 (ONNX) + head
def embed_queries(seqs, order, meta, deadline, sess=None):
    """Mean-pooled ESM embeddings (raw, not normalised) for queries in `order` until `deadline`. -> {query index: vector}"""
    if sess is None:
        import onnxruntime as ort
        so = ort.SessionOptions(); so.intra_op_num_threads = 2; so.inter_op_num_threads = 1
        sess = ort.InferenceSession(str(WDIR / 'esm' / 'model.onnx'), so, providers=['CPUExecutionProvider'])
    v = meta['vocab']; unk, cls, eos, pad = v['<unk>'], v['<cls>'], v['<eos>'], v['<pad>']
    maxlen, budget = meta['maxlen'], meta['tok_budget']
    out = {}
    for st in range(0, len(order), 96):
        if time.time() > deadline: break
        blk = sorted(order[st:st + 96], key=lambda i: len(seqs[i]))
        i = 0
        while i < len(blk):
            j = i + 1
            while j < len(blk) and (j - i + 1) * (min(len(seqs[blk[j]]), maxlen) + 2) <= budget: j += 1
            lst = [[cls] + [v.get(c, unk) for c in seqs[k][:maxlen]] + [eos] for k in blk[i:j]]
            L = max(len(x) for x in lst); ids = np.full((len(lst), L), pad, np.int64); am = np.zeros((len(lst), L), np.int64)
            for r, x in enumerate(lst): ids[r, :len(x)] = x; am[r, :len(x)] = 1
            hid = sess.run(None, {'input_ids': ids, 'attention_mask': am})[0]
            for r, x in enumerate(lst):
                if len(x) > 2: out[blk[i + r]] = hid[r, 1:len(x) - 1].mean(0)
            i = j
    return out


def norm_emb(E, hw):
    X = E.astype(np.float32) - hw['MU']; return X / np.maximum(np.linalg.norm(X, axis=1, keepdims=True), 1e-8)


def knn_search(E, keys, hw, path, K=20):
    """Cosine nearest neighbours of embedded queries among all training embeddings. -> {query index: (positions, sims)}"""
    out = {}
    try:
        T = np.load(path).astype(np.float32)
        for st in range(0, len(keys), 256):
            if time.time() > T0 + KNN_DEADLINE_S: log(f'kNN deadline: {st}/{len(keys)} done'); break
            ks = keys[st:st + 256]; S = norm_emb(np.stack([E[k] for k in ks]), hw) @ T.T
            ix = np.argpartition(-S, K, axis=1)[:, :K]; sv = np.take_along_axis(S, ix, 1); o = np.argsort(-sv, axis=1)
            for r, k in enumerate(ks): out[k] = (ix[r][o[r]], sv[r][o[r]])
    except Exception as e:
        log('kNN failed, continuing without it:', repr(e)[:400]); out = {}
    return out


def head2_and_protos(E, hw2, pw, ctx):
    """-> (within-family probabilities, prototype cosine similarities), each (n, n_ECs)."""
    X = norm_emb(E, hw2)
    Z = ecpipe.gelu(X @ hw2['W1'].T + hw2['b1']) @ hw2['W2'].T + hw2['b2']
    H2 = ecpipe.family_softmax(Z, ctx.l3_id, ctx.nl3).astype(np.float32)
    Y = ecpipe.gelu(X @ pw['W1'].T + pw['b1']) @ pw['W2'].T + pw['b2']
    Y = Y / np.maximum(np.linalg.norm(Y, axis=1, keepdims=True), 1e-8)
    PS = Y @ pw['P'].astype(np.float32).T; PS[:, ~pw['has']] = np.nan
    return H2, PS


def head_probs(E, hw):
    X = norm_emb(E, hw)
    return ecpipe.softmax(ecpipe.gelu(X @ hw['W1'].T + hw['b1']) @ hw['W2'].T + hw['b2'])


# ---------------------------------------------------------------- main
def main(input_path, output_path):
    from inputs import prepare_inputs
    diagnostics()
    meta = json.load(open(WDIR / 'meta.json'))
    raw = pd.read_parquet(input_path)
    frame = prepare_inputs(raw)
    if 'sequence' not in frame.columns:
        frame = frame.merge(raw[['id', 'sequence']].drop_duplicates('id'), on='id')
    ids, seqs = frame.id.tolist(), frame.sequence.tolist(); n = len(seqs)
    log(f'{n} unique queries')
    with gzip.open(WDIR / 'labels.pkl.gz', 'rb') as f: d = pickle.load(f)
    ctx = ecpipe.Ctx(d['labels'], d['ecs'], d['prior'])
    ctx.use_unseen = bool(meta.get('use_unseen')); log(f'unseen-EC candidates: {ctx.use_unseen} ({sum(len(v) for v in ctx.unseen_by_l3.values())} vocabulary ECs without training sequences)')
    if meta.get('use_nb'):
        z = np.load(WDIR / 'nb.npz'); ctx.set_nb(z['E'], z['S']); log(f'neighbourhood table {z["E"].shape}')
    sj = json.load(open(WDIR / 'stacker.json'))
    forests = [ecpipe.Forest(t) for t in sj['forests']] if isinstance(sj, dict) else [ecpipe.Forest(sj)]      # an ensemble: the scores are averaged

    hits = mm_hits(seqs, meta, deadline=T0 + meta.get('mm_deadline_s', MM_DEADLINE_S), chunk=(MM_CHUNK if meta['mm_max_seqs'] <= 100 else 200), budget=meta.get('mm_budget_s', 230))
    log(f'mmseqs hits for {len(hits)}/{n} queries')
    m2 = meta.get('mm2')
    if m2:   # second, more sensitive pass only for queries without a convincing first-pass hit
        weak = [i for i in range(n) if (hits[i][1][0] if i in hits else 0.0) < m2['tb']]
        log(f'second pass for {len(weak)} weak queries (s={m2["s"]}, max-seqs {m2["max_seqs"]})')
        h2 = mm_hits(seqs, meta, T0 + meta.get('mm2_deadline_s', MM2_DEADLINE_S), idxs=weak, s=m2['s'], max_seqs=m2['max_seqs'], chunk=m2.get('chunk', 200), tag='mmseqs2')
        hits.update(h2); log(f'second pass gave hits for {len(h2)}/{len(weak)}; total hits {len(hits)}/{n}')

    H, H2, PS = {}, {}, {}
    if meta.get('use_plm') and (WDIR / 'esm' / 'model.onnx').exists():
        try:
            topbits = np.array([hits[i][1][0] if i in hits else 0.0 for i in range(n)])
            order = [int(i) for i in np.argsort(topbits, kind='stable') if topbits[i] < meta['embed_max_topbits']]
            E = embed_queries(seqs, order, meta, deadline=T0 + EMBED_DEADLINE_S)
            log(f'embedded {len(E)}/{len(order)} candidate queries')
            hw = dict(np.load(WDIR / 'head.npz'))
            keys = list(E)
            for st in range(0, len(keys), 500):
                ks = keys[st:st + 500]; P = head_probs(np.stack([E[k] for k in ks]), hw)
                for k, p in zip(ks, P): H[k] = p
            if meta.get('use_hier'):
                hw2 = dict(np.load(WDIR / 'head2.npz')); pw = dict(np.load(WDIR / 'proto.npz')); pw['has'] = pw['has'].astype(bool)
                for st in range(0, len(keys), 500):
                    ks = keys[st:st + 500]; A2, A3 = head2_and_protos(np.stack([E[k] for k in ks]), hw2, pw, ctx)
                    for k, a2, a3 in zip(ks, A2, A3): H2[k] = a2; PS[k] = a3
                ctx.use_hier_cands = bool(meta.get('use_hier_cands'))
        except Exception as e:
            log('PLM path failed, continuing without it:', repr(e)[:500]); H = {}; H2 = {}; PS = {}; E = {}; hw = None
    log(f'head probabilities for {len(H)} queries | within-family head and prototypes for {len(H2)}')
    KNN = {}
    if H and (WDIR / 'knn_emb.npy').exists():
        KNN = knn_search(E, list(H), hw, WDIR / 'knn_emb.npy'); log(f'kNN for {len(KNN)} queries')

    MOT = {}
    if meta.get('use_motif') and (WDIR / 'motif.npz').exists():
        try:
            mfs = meta.get('motif_files') or [dict(file='motif.npz', ks=[3, 4, 5, 6])]          # several motif models: probabilities are averaged
            mws = [(dict(np.load(WDIR / m['file'])), tuple(m['ks'])) for m in mfs]; t_m = time.time()
            for st in range(0, n, 1000):
                if time.time() - T0 > MOTIF_DEADLINE_S: log(f'motif deadline: stopping at {st}/{n}'); break
                P = np.mean([ecpipe.motif_probs(seqs[st:st + 1000], mw, ks) for mw, ks in mws], axis=0)
                for k, p in enumerate(P): MOT[st + k] = p
            ctx.use_motif_cands = bool(meta.get('use_motif_cands'))
            log(f'motif scores for {len(MOT)}/{n} queries in {time.time() - t_m:.0f}s')
        except Exception as e:
            log('motif path failed, continuing without it:', repr(e)[:500]); MOT = {}
    rows = []
    n2c = 3 if meta.get('nb_cands', True) else 0          # neighbourhood features with or without the two-hop candidates
    for i in range(n): rows += ecpipe.query_rows(i, hits.get(i), len(seqs[i]), H.get(i), ctx, knn=KNN.get(i), h2=H2.get(i), ps=PS.get(i), mot=MOT.get(i), TOPN2C=n2c)
    log(f'{len(rows)} candidate rows')
    best = {}
    if rows:
        A = np.array(rows, dtype=np.float64)
        cols = [ecpipe.COLS.index(f) for f in meta['feats']]; pred = np.mean([f.predict(A[:, cols]) for f in forests], axis=0)
        df = pd.DataFrame(dict(q=A[:, 0].astype(int), e=A[:, 1].astype(int), p=pred))
        best = df.loc[df.groupby('q').p.idxmax()].set_index('q').e.to_dict()
    vp = ROOT / 'ec_vocabulary.json'
    vocab = set(json.loads(vp.read_text())) if vp.exists() else None
    with Path(output_path).open('w') as f:
        for i, qid in enumerate(ids):
            ec = ctx.ecs[int(best[i])] if i in best else ctx.top_ec
            if vocab is not None and ec not in vocab: ec = ctx.top_ec if ctx.top_ec in vocab else sorted(vocab)[0]
            f.write(json.dumps({'id': qid, 'prediction': [ec]}) + '\n')
    log('fallback (no candidates) for', sum(1 for i in range(n) if i not in best), 'of', n, 'queries; done')


if __name__ == '__main__':
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument('--input', required=True); ap.add_argument('--output', required=True)
    a = ap.parse_args(); main(a.input, a.output)