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 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())

    df = pd.read_parquet(args.input, columns=['id', 'sequence'])
    # identical sequences count once; the first existing ID is kept (leaderboard rule)
    uniq = df.drop_duplicates('sequence', keep='first').reset_index(drop=True)
    qid = ['q%d' % i for i in range(len(uniq))]

    with tempfile.TemporaryDirectory(dir='/tmp') as tmpdir:
        tmp = Path(tmpdir)
        with open(tmp / 'q.fa', 'w') as f:
            for i, s in zip(qid, uniq['sequence']):
                f.write('>%s\n%s\n' % (i, s))
        out = tmp / 'hits.tsv'
        cmd = [str(HERE / 'bin' / 'mmseqs'), 'easy-search', str(tmp / 'q.fa'), str(HERE / 'db' / 'train.fa'),
               str(out), str(tmp / 'work'), '-s', str(cfg['sens']), '--max-seqs', '100', '-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:
            hits = pd.read_csv(out, sep='\t', names=['query', 'target', 'bits'])
        else:
            hits = pd.DataFrame(columns=['query', 'target', 'bits'])
    sys.stderr.write('search done after %.0fs\n' % (time.perf_counter() - t0))

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


if __name__ == '__main__':
    main()