"""Offline CPU inference: ensemble of fine-tuned reaction BERTs (int8) blended with a TF-IDF reaction kNN.
Output: top-5 distinct ECs per unique input, best first (only the first is scored)."""
import argparse
import json
import re
from pathlib import Path

import numpy as np
import pandas as pd
import scipy.sparse as sp
import torch
import torch.nn as nn
from transformers import BertConfig, BertModel

import featurize as F
from inputs import prepare_inputs

SMI_RE = re.compile(r"(\[[^\]]+]|Br?|Cl?|N|O|S|P|F|I|b|c|n|o|s|p|\(|\)|\.|=|#|-|\+|\\|/|:|~|@|\?|>>?|\*|\$|%[0-9]{2}|[0-9])")
MAX_LEN = 256
TOP = 5


class Net(nn.Module):
    def __init__(self, config, n_cls, n_aux):
        super().__init__()
        self.bert = BertModel(config, add_pooling_layer=False)
        self.head = nn.Linear(config.hidden_size, n_cls)
        self.aux = nn.ModuleList([nn.Linear(config.hidden_size, n) for n in n_aux])

    def forward(self, ids, mask):
        x = self.bert(input_ids=ids, attention_mask=mask).last_hidden_state
        m = mask.unsqueeze(-1).to(x.dtype)
        return self.head((x * m).sum(1) / m.sum(1))


def bert_probs(wdir, rxns, all_idx, quantize):
    """Softmax probabilities over the shared class list (all_idx maps this model's classes into it)."""
    classes = json.loads((wdir / 'classes.json').read_text())
    tokens = (wdir / 'vocab.txt').read_text().splitlines()
    idx = {t: i for i, t in enumerate(tokens)}
    state = torch.load(wdir / 'best.pt', map_location='cpu', weights_only=True)
    n_aux = [state[f'aux.{i}.weight'].shape[0] for i in range(3)] if 'aux.0.weight' in state else []
    net = Net(BertConfig.from_pretrained(wdir), len(classes), n_aux)
    net.load_state_dict(state, strict=False)
    net.eval()
    if quantize:
        net = torch.quantization.quantize_dynamic(net, {nn.Linear}, dtype=torch.qint8)
    enc = [[idx['[CLS]']] + [idx.get(t, 1) for t in SMI_RE.findall(r)][:MAX_LEN - 2] + [idx['[SEP]']] for r in rxns]
    order = sorted(range(len(enc)), key=lambda i: len(enc[i]))
    out = np.zeros((len(rxns), len(all_idx[1])), np.float32)
    with torch.inference_mode():
        for a in range(0, len(order), 64):
            b = order[a:a + 64]
            n = max(len(enc[i]) for i in b)
            ids = torch.zeros(len(b), n, dtype=torch.long)
            mask = torch.zeros(len(b), n, dtype=torch.long)
            for j, i in enumerate(b):
                ids[j, :len(enc[i])] = torch.tensor(enc[i]); mask[j, :len(enc[i])] = 1
            p = torch.softmax(net(ids, mask).float(), -1).numpy()
            for j, i in enumerate(b):
                out[i, all_idx[0]] = p[j]
    return out


def knn_probs(kdir, Q, n, all_cls, k=10, power=8):
    X = sp.load_npz(kdir / 'X.npz').tocsr()
    Y = sp.load_npz(kdir / 'Y.npz').tocsr()
    kc = json.loads((kdir / 'classes.json').read_text())
    XT = X.T.tocsr()
    out = np.zeros((n, len(kc)), np.float32)
    for a in range(0, n, 250):
        S = (Q[a:a + 250] @ XT).toarray()
        ik = np.argpartition(-S, k, axis=1)[:, :k]
        sk = np.maximum(np.take_along_axis(S, ik, 1), 0) ** power
        for j in range(k):
            out[a:a + 250] += sk[:, j:j + 1] * Y[ik[:, j]].toarray()
    out /= np.maximum(out.sum(1, keepdims=True), 1e-9)
    pos = {e: i for i, e in enumerate(all_cls)}
    full = np.zeros((n, len(all_cls)), np.float32)
    for i, e in enumerate(kc):
        if e in pos:
            full[:, pos[e]] = out[:, i]
    return full


def mlp_probs(mdir, Q, all_cls):
    """Fingerprint MLP (Linear-ReLU-Linear) on the same TF-IDF features as the kNN."""
    st = torch.load(mdir / 'mlp.pt', map_location='cpu', weights_only=True)
    cl = json.loads((mdir / 'classes.json').read_text())
    net = nn.Sequential(nn.Linear(st['0.weight'].shape[1], st['0.weight'].shape[0]), nn.ReLU(), nn.Dropout(0.0),
                        nn.Linear(st['3.weight'].shape[1], st['3.weight'].shape[0]))
    net.load_state_dict({k: v.float() for k, v in st.items()})
    net.eval()
    pos = {e: i for i, e in enumerate(all_cls)}
    idx = [pos[e] for e in cl]
    full = np.zeros((Q.shape[0], len(all_cls)), np.float32)
    with torch.inference_mode():
        for a in range(0, Q.shape[0], 512):
            full[a:a + 512][:, idx] = torch.softmax(net(torch.from_numpy(Q[a:a + 512].toarray())), -1).numpy()
    return full


def predict_file(input_path, output_path):
    torch.set_num_threads(2)
    root = Path(__file__).parent
    w = root / 'weights'
    cfg = json.loads((w / 'blend.json').read_text())
    frame = prepare_inputs(pd.read_parquet(input_path))
    rxns = frame.rxn_smiles.tolist()
    valid = set(json.loads((root / 'ec_vocabulary.json').read_text()))
    kdir = w / 'knn'
    all_cls = json.loads((kdir / 'classes.json').read_text())
    pos = {e: i for i, e in enumerate(all_cls)}
    probs = 0
    models = sorted(d for d in w.iterdir() if d.is_dir() and (d / 'best.pt').exists())
    for d in models:
        cl = json.loads((d / 'classes.json').read_text())
        probs = probs + bert_probs(d, rxns, ([pos[e] for e in cl], all_cls), cfg['quantize'])
    probs = probs / len(models)
    kw, mw = cfg['knn_weight'], cfg.get('mlp_weight', 0.0)
    idf = np.load(kdir / 'idf.npz')
    whole, diff = F.raw_features(rxns)
    Q = F.combine(whole, diff, idf['whole'], idf['diff'])
    blend = (1 - kw - mw) * probs + kw * knn_probs(kdir, Q, len(rxns), all_cls)
    if mw:
        blend = blend + mw * mlp_probs(w / 'mlp', Q, all_cls)
    probs = blend
    top = np.argsort(-probs, axis=1)[:, :TOP]
    with Path(output_path).open('w') as out:
        for key, row in zip(frame.id, top):
            p = [all_cls[t] for t in row if all_cls[t] in valid]
            out.write(json.dumps({'id': key, 'prediction': p or [all_cls[0]]}) + '\n')


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