"""Feature extraction shared by train.py and model.py (must be identical at train and inference time)."""
from itertools import product

import numpy as np
import pandas as pd
from rdkit import Chem, RDLogger
from rdkit.Chem import Descriptors, MACCSkeys, rdFingerprintGenerator

RDLogger.DisableLog('rdApp.*')
TARGETS = ['log10_kcat', 'log10_km']

# ---------------- enzyme: handcrafted composition features (489) ----------------
AA = 'ACDEFGHIKLMNPQRSTVWY'
DIP = {a + b: i for i, (a, b) in enumerate(product(AA, AA))}
KD = dict(A=1.8, R=-4.5, N=-3.5, D=-3.5, C=2.5, Q=-3.5, E=-3.5, G=-0.4, H=-3.2, I=4.5, L=3.8, K=-3.9,
          M=1.9, F=2.8, P=-1.6, S=-0.8, T=-0.7, W=-0.9, Y=-1.3, V=4.2)
AA_GROUPS = ['AVLIMFWC', 'STNQGPY', 'KRH', 'DE', 'FWYH', 'AGSCTP']


def seq_feats(seq):
    s = ''.join(c for c in seq.upper() if c in AA) or 'A'
    L = len(s)
    aac = np.array([s.count(a) for a in AA]) / L
    dc = np.zeros(400)
    for i in range(L - 1):
        dc[DIP[s[i:i + 2]]] += 1
    dc /= max(1, L - 1)
    kd = np.array([KD[c] for c in s])
    grp = [sum(s.count(c) for c in g) / L for g in AA_GROUPS]
    thirds = [np.array([p.count(a) for a in AA]) / max(1, len(p))
              for p in (s[:L // 3], s[L // 3:2 * L // 3], s[2 * L // 3:])]
    return np.concatenate([[np.log(L), kd.mean(), kd.std()], grp, aac, dc, *thirds]).astype(np.float32)


# ---------------- substrate: RDKit features (2048 Morgan counts + 167 MACCS + descriptors) ----------------
_MORGAN = rdFingerprintGenerator.GetMorganGenerator(radius=2, fpSize=2048)
N_SMI = 2048 + 167 + len(Descriptors._descList)


def smi_feats(smi):
    m = Chem.MolFromSmiles(smi)
    if m is None:
        return np.full(N_SMI, np.nan, dtype=np.float32)
    fp = np.asarray(_MORGAN.GetCountFingerprintAsNumPy(m), dtype=np.float32)
    mac = np.array(list(MACCSkeys.GenMACCSKeys(m)), dtype=np.float32)
    d = []
    for _, f in Descriptors._descList:
        try:
            d.append(f(m))
        except Exception:
            d.append(np.nan)
    d = np.nan_to_num(np.array(d, dtype=np.float64), nan=np.nan, posinf=1e6, neginf=-1e6).clip(-1e6, 1e6)
    return np.concatenate([fp, mac, d]).astype(np.float32)


def base_features(sequences, smiles):
    """Enzyme + substrate features, computed once per unique string."""
    S = {s: seq_feats(s) for s in pd.unique(sequences)}
    M = {s: smi_feats(s) for s in pd.unique(smiles)}
    return np.hstack([np.stack([S[s] for s in sequences]), np.stack([M[s] for s in smiles])])


# ---------------- substrate prior: statistics of each substrate's labels in the training data ----------------
PRIOR_K = 5  # shrinkage: (n*mean + K*global_mean) / (n + K)


def prior_table(smiles, Y):
    """Per-substrate label statistics from training rows -> DataFrame indexed by SMILES."""
    d = pd.DataFrame(Y, columns=TARGETS)
    g = d.groupby(np.asarray(smiles))
    m, n = g.mean(), g.size()
    gm = Y.mean(0)
    tab = (m.mul(n, axis=0) + gm * PRIOR_K).div(n + PRIOR_K, axis=0).add_prefix('shrunk_')
    tab['log_count'] = np.log1p(n)
    for agg in ['std', 'max', 'min', 'median']:
        tab = tab.join(getattr(g, agg)().add_prefix(agg + '_'))
    tab.attrs['global_mean'] = gm.tolist()
    return tab


def prior_features(tab, smiles, global_mean):
    """Look up prior features; unseen substrates get the global mean, count 0 and NaN spread stats."""
    out = tab.reindex(np.asarray(smiles)).to_numpy(dtype=np.float64).copy()
    unseen = np.isnan(out[:, 0])
    out[unseen, 0], out[unseen, 1], out[unseen, 2] = global_mean[0], global_mean[1], 0.0
    return out


# ---------------- enzyme: ESM-2 protein language model embeddings (optional, v2) ----------------
def load_esm(model_dir, device='cpu', quantize=False, layer=None):
    """layer=k keeps only the first k transformer layers and drops the final LayerNorm, so last_hidden_state equals
    hidden_states[k] of the full model (an intermediate layer) and inference is faster."""
    import torch
    from transformers import AutoTokenizer, EsmModel
    tok = AutoTokenizer.from_pretrained(model_dir)
    model = EsmModel.from_pretrained(model_dir, add_pooling_layer=False)
    if layer:
        model.encoder.layer = model.encoder.layer[:layer]
        model.encoder.emb_layer_norm_after = None
        model.contact_head = None   # unused for embeddings; its size depends on the layer count
        model.config.num_hidden_layers = layer
    model = model.to(device).eval()
    if quantize:  # int8 weights for Linear layers: ~1.6x faster on CPU, embeddings ~0.999 cosine to fp32
        model = torch.ao.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)
    return tok, model


def esm_embed(seqs, tok, model, device='cpu', batch_tokens=2000, window=1022):
    """Mean-pooled last-layer embedding per sequence (special tokens excluded).
    Sequences longer than `window` are split into windows whose embeddings are averaged."""
    import torch
    pieces, owner = [], []
    for i, s in enumerate(seqs):
        s = s.upper()
        for j in range(0, max(1, len(s)), window):
            pieces.append(s[j:j + window]); owner.append(i)
    order = np.argsort([len(p) for p in pieces])[::-1]
    out = np.zeros((len(pieces), model.config.hidden_size), dtype=np.float32)
    k = 0
    with torch.no_grad():
        while k < len(order):
            n = max(1, batch_tokens // (len(pieces[order[k]]) + 2))
            idx = order[k:k + n]; k += n
            enc = tok([pieces[i] for i in idx], return_tensors='pt', padding=True).to(device)
            h = model(**enc).last_hidden_state.float()
            m = enc['attention_mask'].clone(); m[:, 0] = 0
            m[torch.arange(len(idx)), enc['attention_mask'].sum(1) - 1] = 0
            out[idx] = ((h * m.unsqueeze(-1)).sum(1) / m.sum(1, keepdim=True)).cpu().numpy()
    owner = np.array(owner)
    return np.stack([out[owner == i].mean(0) for i in range(len(seqs))])