import argparse
import json
from pathlib import Path

import numpy as np
import pandas as pd
import torch

import diagindex as di
import dlmodel as dl
import mmsearch as ms
from inputs import prepare_inputs

BETA, GATE, BUDGET, SHARE = 0.7, 20, 420, 0.95


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 = di.candidates(index, di.to_codes(seq))
        out[r] = di.vote(cand, s, lab_flat, lab_ptr, n_ec)
        best[r] = s[0] if len(s) else 0
    return out, best


def predict_file(input_path, output_path):
    torch.set_num_threads(2)
    root = Path(__file__).parent
    frame = prepare_inputs(pd.read_parquet(input_path))
    seqs = frame.sequence.tolist()
    ec_list = json.loads((root / 'weights' / 'meta.json').read_text())['ec_list']
    n_ec = len(ec_list)

    z = np.load(root / 'weights' / 'kmer_index.npz')
    codes, lens = z['codes'], z['lens']
    lab_flat = z['lab_flat'].astype(np.int64)
    lab_ptr = np.concatenate([[0], np.cumsum(z['lab_len'])])

    try:
        hits = ms.search(root, codes, lens, seqs, threads=2, budget=BUDGET)
        score, has = ms.votes(hits, lab_flat, lab_ptr, len(seqs), n_ec)
    except Exception as error:
        print('mmseqs failed, using the fallback for every input:', error, flush=True)
        score, has = np.zeros((len(seqs), n_ec), dtype=np.float32), np.zeros(len(seqs), dtype=bool)

    unsure = np.where(has & (score.max(1) < SHARE))[0]
    nohit = np.where(~has)[0]
    print('unsure inputs:', len(unsure), ' inputs without a search hit:', len(nohit), ' of', len(seqs), flush=True)

    need = np.concatenate([unsure, nohit]).astype(np.int64)
    if len(need):
        index = di.build(codes, lens)
        red, best = red_votes(index, lab_flat, lab_ptr, [seqs[i] for i in need], n_ec)
        score[unsure] += red[:len(unsure)]
        if len(nohit):
            red_n, best_n = red[len(unsure):], best[len(unsure):]
            weak = np.where(best_n < GATE)[0]
            if len(weak):
                P, cfg, bank, alpha = dl.load(root / 'weights')
                emb, probs = dl.run(P, cfg, [seqs[nohit[i]] for i in weak], 'cpu')
                cnn = alpha * dl.knn_votes(bank, emb) + (1 - alpha) * probs
                red_n[weak] = BETA * red_n[weak] + (1 - BETA) * cnn
            score[nohit] = red_n

    preds = [ec_list[i] for i in score.argmax(axis=1)]
    valid = set(json.loads((root / 'ec_vocabulary.json').read_text()))
    assert set(preds) <= valid, 'prediction outside ec_vocabulary.json'
    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')


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)