import argparse
import json
import subprocess
import sys
import tempfile
import time
from pathlib import Path

import pandas as pd

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
import knn


def search_chunk(rows, seqs, cfg):
    with tempfile.TemporaryDirectory(dir='/tmp') as td:
        td = Path(td)
        with open(td / 'q.fa', 'w') as f:
            for i in rows:
                f.write('>q%d\n%s\n' % (i, seqs[i]))
        out = td / 'hits.tsv'
        cmd = [str(HERE / 'bin' / 'mmseqs'), 'easy-search', str(td / 'q.fa'), str(HERE / 'db' / 'trainDB'), str(out),
               str(td / 'work'), '-s', str(cfg['sens']), '--max-seqs', str(cfg['max_seqs']), '-e', '10',
               '--split-memory-limit', cfg['split_mem'], '--threads', str(cfg['threads']),
               '--format-output', 'query,target,bits', '-v', '1']
        r = subprocess.run(cmd, capture_output=True, text=True)
        if r.returncode != 0:
            sys.stderr.write(r.stdout[-2000:] + '\n' + r.stderr[-2000:] + '\n')
            sys.exit('mmseqs failed with code %d' % r.returncode)
        if out.exists() and out.stat().st_size > 0:
            return pd.read_csv(out, sep='\t', names=['query', 'target', 'bits'])
        return pd.DataFrame(columns=['query', 'target', 'bits'])


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--input', required=True)
    ap.add_argument('--output', required=True)
    args = ap.parse_args()
    t0 = time.perf_counter()
    cfg = json.loads((HERE / 'config.json').read_text())
    ec_map = json.loads((HERE / 'db' / 'ec_map.json').read_text())
    fallback = json.loads((HERE / 'db' / 'fallback.json').read_text())
    freq = json.loads((HERE / 'db' / 'freq.json').read_text())

    df = pd.read_parquet(args.input, columns=['id', 'sequence'])
    uniq = df.drop_duplicates('sequence', keep='first').reset_index(drop=True)   # identical inputs count once, first ID kept
    seqs = uniq['sequence'].tolist()
    order = sorted(range(len(seqs)), key=lambda i: len(seqs[i]))                  # short first: most queries finish if time runs out
    chunks = [order[i:i + cfg['chunk']] for i in range(0, len(order), cfg['chunk'])]

    hl, times, skipped = {}, [], 0
    for ci, rows in enumerate(chunks):
        elapsed = time.perf_counter() - t0
        if times and elapsed + 1.3 * max(times) > cfg['time_budget']:
            skipped = sum(len(c) for c in chunks[ci:])
            sys.stderr.write('TIME GUARD: stopping before chunk %d; %d inputs get the fallback EC\n' % (ci, skipped))
            break
        tc = time.perf_counter()
        hits = search_chunk(rows, seqs, cfg)
        hl.update(knn.hits_from_df(hits, cfg['kmax']))
        times.append(time.perf_counter() - tc)
        sys.stderr.write('chunk %d/%d done in %.0fs (elapsed %.0fs)\n' % (ci + 1, len(chunks), times[-1], time.perf_counter() - t0))

    with open(args.output, 'w') as f:
        for i, rid in enumerate(uniq['id']):
            ecs = knn.select_ecs(hl.get('q%d' % i, []), ec_map, fallback, cfg['K'], cfg['pw'], cfg['ratio'], cfg['max_ec'], freq)
            f.write(json.dumps({'id': rid, 'prediction': ecs}) + '\n')
    sys.stderr.write('total %.0fs for %d unique inputs (%d without search)\n' % (time.perf_counter() - t0, len(seqs), skipped))


if __name__ == '__main__':
    main()