"""Offline CPU inference: k-mer prefilter (5-mers + 4-mers) -> Smith-Waterman re-ranking -> top-20 weighted EC vote."""
import argparse
import gc
import json
import os
import time
from collections import Counter, defaultdict
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import numpy as np
import pandas as pd
import pyopal

from inputs import prepare_inputs
from kmer import KmerIndex

ROOT = Path(__file__).parent
T0 = time.time()
BUDGET = float(os.environ.get('TIME_BUDGET', 480))  # seconds; hard limit is 600
N5 = 1000          # 5-mer candidates per query
N4 = 300           # extra 4-mer candidates per query (merged, duplicates removed)
TOP_K = 20         # neighbours that vote
EXPONENT = 16      # vote weight = (score / best score) ** EXPONENT
MAX_LEN_RATIO = 3  # skip 4-mer extras longer than 3x the query (slow to align, rarely homologs)
THREADS = 2


def log(msg):
    print(f'[{time.time() - T0:6.1f}s] {msg}', flush=True)


def load_db():
    path = ROOT / 'weights/db.parquet'
    if not path.exists():
        path = ROOT / 'db.parquet'
    db = pd.read_parquet(path)
    vocab = set(json.loads((ROOT / 'ec_vocabulary.json').read_text()))
    ecs = [[e for e in row if e in vocab] for row in db.ecs]
    keep = [i for i, e in enumerate(ecs) if e]
    seqs = [db.sequence.iat[i] for i in keep]
    ecs = [ecs[i] for i in keep]
    fallback = Counter(e for row in ecs for e in row).most_common(1)[0][0]
    return seqs, ecs, fallback


def vote(cand, scores, ecs, fallback):
    if len(cand) == 0 or scores.max() <= 0:
        return fallback
    order = np.argsort(-scores)[:TOP_K]
    best = scores[order[0]]
    votes = defaultdict(float)
    for i in order:
        weight = (max(scores[i], 0) / best) ** EXPONENT
        row = ecs[cand[i]]
        for e in row:
            votes[e] += weight / len(row)
    return max(votes, key=votes.get)


def candidates(seqs, queries, k, n):
    index = KmerIndex(seqs, k=k)
    idx, _ = index.topk(queries, n)
    del index
    gc.collect()
    log(f'{k}-mer candidates ready')
    return idx


def predict_file(input_path, output_path):
    frame = prepare_inputs(pd.read_parquet(input_path))
    queries = [s if isinstance(s, str) else '' for s in frame.sequence]
    log(f'{len(queries)} unique inputs')

    seqs, ecs, fallback = load_db()
    c5 = candidates(seqs, queries, 5, N5)
    # Default prediction: 5-mer nearest neighbour (used if time runs short).
    preds = [ecs[i[0]][0] if i[0] >= 0 else fallback for i in c5]
    c4 = candidates(seqs, queries, 4, N4) if time.time() - T0 < BUDGET / 3 else None

    aligner = pyopal.Aligner('BLOSUM62', gap_open=11, gap_extend=1)
    database = pyopal.Database(seqs)
    lengths = np.array([len(s) for s in seqs])

    def align(q):
        if time.time() - T0 > BUDGET or not queries[q]:
            return q, None
        pool = c5[q]
        if c4 is not None:
            extra = c4[q][c4[q] >= 0]
            extra = extra[lengths[extra] <= MAX_LEN_RATIO * len(queries[q])]
            pool = np.concatenate([pool, extra])
        cand = [int(c) for c in dict.fromkeys(pool.tolist()) if c >= 0]
        if not cand:
            return q, None
        scores = np.array([r.score for r in aligner.align(queries[q], database.extract(cand))], float)
        return q, vote(cand, scores, ecs, fallback)

    done = 0
    with ThreadPoolExecutor(THREADS) as pool:
        for q, p in pool.map(align, range(len(queries))):
            if p is not None:
                preds[q] = p
                done += 1
    log(f'aligned {done}/{len(queries)} queries')

    with Path(output_path).open('w') as out:
        for key, p in zip(frame.id, preds):
            out.write(json.dumps({'id': key, 'prediction': [p]}) + '\n')
    log(f'wrote {output_path}')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--input', required=True)
    parser.add_argument('--output', required=True)
    args = parser.parse_args()
    predict_file(args.input, args.output)