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
import rerank


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', ','.join(rerank.HIT_FIELDS), '-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:
            h = pd.read_csv(out, sep='\t', names=rerank.HIT_FIELDS)
            if h['pident'].max() > 1.0:
                h['pident'] = h['pident'] / 100
            return h
        return pd.DataFrame(columns=rerank.HIT_FIELDS)


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())
    l3_tot = rerank.l3_totals(freq)

    model = None
    if cfg.get('use_rerank'):
        try:
            import joblib
            model = joblib.load(HERE / 'model' / 'rr.joblib')
        except Exception as e:                      # never fail because of the re-ranker: fall back to the vote
            sys.stderr.write('re-ranker unavailable, using the vote only: %r\n' % (e,))

    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 inputs finish if time runs out
    chunks = [order[i:i + cfg['chunk']] for i in range(0, len(order), cfg['chunk'])]

    hl, rr, 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[['query', 'target', 'bits']], cfg['kmax']))
        if model is not None and (time.perf_counter() - t0) < 0.45 * cfg['time_budget']:
            try:
                cand = rerank.build_candidates(hits, ec_map, freq, cfg['kmax'], l3_tot)
                rr.update(rerank.rerank_select(cand, model, cfg['rel'], cfg['max_ec']))
            except Exception as e:
                sys.stderr.write('re-ranker failed, using the vote for the rest: %r\n' % (e,))
                model = None
        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))

    changed = 0
    with open(args.output, 'w') as f:
        for i, rid in enumerate(uniq['id']):
            q = 'q%d' % i
            vote = knn.select_ecs(hl.get(q, []), ec_map, fallback, cfg['K'], cfg['pw'], cfg['ratio'], cfg['max_ec'], freq)
            ecs = rr.get(q) or vote
            changed += int(ecs[0] != vote[0])
            f.write(json.dumps({'id': rid, 'prediction': ecs}) + '\n')
    sys.stderr.write('re-ranker used: %s | first EC differs from the vote for %d inputs\n' % (bool(rr), changed))
    sys.stderr.write('total %.0fs for %d unique inputs (%d without search)\n' % (time.perf_counter() - t0, len(seqs), skipped))


if __name__ == '__main__':
    main()