"""Homology-graph neural network for EC prediction (offline, 2 CPUs).

1. Retrieval: k-mer prefilter (5-mers + 4-mers) finds candidate training enzymes; Smith-Waterman alignment
   (BLOSUM62) scores them; the 20 best hits form the query's homology graph.
2. Neural network: a graph attention network (gnn.py) reads the 20 neighbours (alignment features + EC agreement),
   passes messages between them with two transformer layers and assigns each neighbour a weight.
   The EC carrying the most weight is the prediction.
"""
import argparse
import gc
import json
import os
import time
from collections import Counter
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import numpy as np
import pandas as pd
import pyopal
import scoring_matrices

from gnn import K, HomologyGNN, features
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)
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 find(name):
    for p in (ROOT / 'weights' / name, ROOT / name):
        if p.exists():
            return p
    raise FileNotFoundError(name)


def load_db():
    db = pd.read_parquet(find('db.parquet'))
    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 self_scores(seqs):
    m = scoring_matrices.ScoringMatrix.from_name('BLOSUM62')
    diag = np.zeros(256)
    for i, ch in enumerate(m.alphabet):
        diag[ord(ch)] = m[ch][i]
    return np.array([diag[np.frombuffer(s.encode('ascii', 'replace'), np.uint8)].sum() for s in seqs], np.float32)


def prefix_ids(ecs):
    """(N,4) ids of the L1..L4 prefixes of each enzyme's first EC."""
    pid = {}
    return np.array([[pid.setdefault('.'.join(row[0].split('.')[:L]) + f'#{L}', len(pid)) for L in (1, 2, 3, 4)]
                     for row in ecs], np.int64)


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 only 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 = np.array([int(c) for c in dict.fromkeys(pool.tolist()) if c >= 0])
        if len(cand) == 0:
            return q, None
        scores = np.array([r.score for r in aligner.align(queries[q], database.extract(cand.tolist()))], np.float32)
        if scores.max() <= 0:
            return q, None
        o = np.argsort(-scores)[:K]
        return q, (cand[o], scores[o])

    NI = np.full((len(queries), K), -1, np.int64)
    NS = np.zeros((len(queries), K), np.float32)
    with ThreadPoolExecutor(THREADS) as pool:
        for q, res in pool.map(align, range(len(queries))):
            if res is not None:
                NI[q, :len(res[0])], NS[q, :len(res[1])] = res
    aligned = np.where(NI[:, 0] >= 0)[0]
    log(f'aligned {len(aligned)}/{len(queries)} queries')

    # Homology graph neural network
    net = HomologyGNN(find('gnn.npz'))
    dself, dpref = self_scores(seqs), prefix_ids(ecs)
    qself = self_scores(queries)
    qlen = np.array([max(len(s), 1) for s in queries], np.float32)
    for s in range(0, len(aligned), 1024):
        b = aligned[s:s + 1024]
        ni, valid = NI[b].clip(0), NI[b] >= 0
        f, rel = features(NS[b], qself[b], dself[ni], qlen[b], lengths[ni].astype(np.float32), valid)
        w = net.neighbour_weights(f, rel, dpref[ni], valid)
        for r, q in enumerate(b):
            mass = {}
            for j in range(K):
                if valid[r, j]:
                    row = ecs[NI[q, j]]
                    for e in row:
                        mass[e] = mass.get(e, 0) + w[r, j] / len(row)
            preds[q] = max(mass, key=mass.get)
    log(f'graph network predicted {len(aligned)} 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)