"""Shared feature code. train.py and predict.py BOTH import this file,
so features are guaranteed to be identical at training and inference time."""
import time
import zlib
from collections import Counter

import numpy as np

AA = "ACDEFGHIKLMNPQRSTVWY"
_LUT = np.full(256, -1, dtype=np.int16)
for _i, _a in enumerate(AA):
    _LUT[ord(_a)] = _i

N_MORGAN = 1024


# ----------------------------------------------------------------- SMILES ---
def _smiles_vec_factory():
    from rdkit import Chem, RDLogger
    from rdkit.Chem import (Descriptors, MACCSkeys, rdFingerprintGenerator,
                            rdMolDescriptors)
    RDLogger.DisableLog("rdApp.*")
    gen = rdFingerprintGenerator.GetMorganGenerator(radius=2, fpSize=N_MORGAN)
    n_desc = 12
    dim = N_MORGAN + 167 + n_desc

    def vec(smi):
        out = np.zeros(dim, dtype=np.float32)
        try:
            mol = Chem.MolFromSmiles(smi)
            if mol is None:
                return out
            out[:N_MORGAN] = gen.GetCountFingerprintAsNumPy(mol)
            out[N_MORGAN:N_MORGAN + 167] = np.array(list(MACCSkeys.GenMACCSKeys(mol)), dtype=np.float32)
            d = [
                Descriptors.MolWt(mol), Descriptors.MolLogP(mol), Descriptors.TPSA(mol),
                Descriptors.NumHDonors(mol), Descriptors.NumHAcceptors(mol),
                Descriptors.NumRotatableBonds(mol), rdMolDescriptors.CalcNumRings(mol),
                rdMolDescriptors.CalcNumAromaticRings(mol), mol.GetNumHeavyAtoms(),
                Chem.GetFormalCharge(mol), rdMolDescriptors.CalcFractionCSP3(mol),
                len(Chem.GetMolFrags(mol)),
            ]
            out[N_MORGAN + 167:] = np.array(d, dtype=np.float32)
        except Exception:
            pass
        return out

    return vec, dim


def smiles_features(unique_smiles):
    """-> (matrix [n_unique, dim]) in the same order as unique_smiles"""
    vec, dim = _smiles_vec_factory()
    return np.stack([vec(s) for s in unique_smiles]).astype(np.float32)


# --------------------------------------------------------------- SEQUENCE ---
def _seq_vec(s):
    a = _LUT[np.frombuffer(s.encode("ascii", "ignore"), dtype=np.uint8)]
    a = a[a >= 0].astype(np.int64)
    n = len(a)
    comp = np.bincount(a, minlength=20) / max(n, 1)
    if n > 1:
        di = np.bincount(a[:-1] * 20 + a[1:], minlength=400) / (n - 1)
    else:
        di = np.zeros(400)
    return np.concatenate([[len(s), np.log1p(len(s))], comp, di]).astype(np.float32)


def seq_features(unique_seqs):
    return np.stack([_seq_vec(s) for s in unique_seqs]).astype(np.float32)


# -------------------------------------------------------------------- ESM ---
def esm_embed(seqs, model_dir, fallback=None, max_len=800, budget_s=None,
              batch_tokens=4096, n_threads=2, device="cpu", log=print, layers=None):
    """Mean-pooled ESM2 embedding for each sequence (list of str).
    Short sequences are done first. If budget_s runs out, the remaining
    sequences get `fallback` (the training-mean embedding) so we never time out."""
    import torch
    from transformers import AutoTokenizer, EsmModel
    torch.set_num_threads(n_threads)
    tok = AutoTokenizer.from_pretrained(model_dir)
    model = EsmModel.from_pretrained(model_dir).to(device).eval()
    dim = model.config.hidden_size * (len(layers) if layers else 1)
    out = np.full((len(seqs), dim), np.nan, dtype=np.float32)
    order = np.argsort([len(s) for s in seqs])
    t0 = time.time()
    i = 0
    while i < len(order):
        if budget_s is not None and time.time() - t0 > budget_s:
            log(f"[esm] time budget hit after {i}/{len(order)} sequences; using fallback for the rest")
            break
        j, maxl = i, 0
        while j < len(order):
            l = min(len(seqs[order[j]]), max_len) + 2
            if j > i and (j - i + 1) * max(maxl, l) > batch_tokens:
                break
            maxl = max(maxl, l)
            j += 1
        idx = order[i:j]
        batch = [seqs[k][:max_len] for k in idx]
        enc = tok(batch, return_tensors="pt", padding=True)
        enc = {k: v.to(device) for k, v in enc.items()}      # move inputs to the same device as the model
        with torch.inference_mode():
            o = model(**enc, output_hidden_states=bool(layers))
        m = enc["attention_mask"].clone().float()
        lens = enc["attention_mask"].sum(1)
        m[:, 0] = 0
        m[torch.arange(len(idx), device=m.device), lens - 1] = 0
        hs = [o.hidden_states[l] for l in layers] if layers else [o.last_hidden_state]
        e = torch.cat([(h * m.unsqueeze(-1)).sum(1) / m.sum(1, keepdim=True).clamp(min=1) for h in hs], dim=1)
        out[idx] = e.cpu().numpy()
        i = j
        if (i // 50) != ((i - len(idx)) // 50):
            log(f"[esm] {i}/{len(order)} done, {time.time() - t0:.0f}s")
    missing = np.isnan(out[:, 0])
    if missing.any():
        out[missing] = fallback if fallback is not None else 0.0
    return out


# ------------------------------------------------------ SEQUENCE CLUSTERS ---
def cluster_sequences(seqs, k=6, mod=2, min_frac=0.3):
    """Cheap grouping of similar sequences (mutants / close homologs).
    Used ONLY to make an honest local validation split. Returns a group id per sequence."""
    order = sorted(range(len(seqs)), key=lambda i: -len(seqs[i]))
    kmer2leaders = {}
    leader_of = [0] * len(seqs)
    for i in order:
        s = seqs[i]
        kms = {s[j:j + k] for j in range(len(s) - k + 1)
               if zlib.crc32(s[j:j + k].encode()) % mod == 0}
        hits = Counter()
        for km in kms:
            for L in kmer2leaders.get(km, ()):
                hits[L] += 1
        best = None
        if hits:
            L, c = hits.most_common(1)[0]
            if c >= min_frac * max(len(kms), 1):
                best = L
        if best is None:
            leader_of[i] = i
            for km in kms:
                kmer2leaders.setdefault(km, []).append(i)
        else:
            leader_of[i] = best
    return np.array(leader_of)


def cluster_by_cosine(seqs, thr=0.35, max_frac=0.2, log=print):
    """Group sequences that are even remotely related (3-mer cosine similarity >= thr),
    so a local holdout looks like the real validation set (whose enzymes are unrelated to training).
    If one giant group appears, the threshold is raised automatically."""
    from sklearn.feature_extraction.text import CountVectorizer
    from sklearn.preprocessing import normalize
    from scipy.sparse import coo_matrix
    from scipy.sparse.csgraph import connected_components
    cv = CountVectorizer(analyzer="char", ngram_range=(3, 3), lowercase=False, dtype=np.float32)
    A = normalize(cv.fit_transform(seqs)).tocsr()
    AT = A.T.tocsr()
    n = A.shape[0]
    rows, cols, vals = [], [], []
    for i in range(0, n, 500):
        S = (A[i:i + 500] @ AT).toarray()
        r, c = np.nonzero(S >= 0.30)
        keep = (r + i) < c
        rows.append(r[keep] + i); cols.append(c[keep]); vals.append(S[r[keep], c[keep]])
    rows, cols, vals = np.concatenate(rows), np.concatenate(cols), np.concatenate(vals)
    t = thr
    while True:
        m = vals >= t
        g = coo_matrix((np.ones(int(m.sum())), (rows[m], cols[m])), shape=(n, n))
        _, lab = connected_components(g, directed=False)
        big = np.bincount(lab).max() / n
        if big <= max_frac or t >= 0.9:
            break
        t += 0.05
    log(f"grouping threshold used: {t:.2f} | groups: {len(set(lab))} | biggest group: {big:.1%} of sequences")
    return lab


def _wmean(Y, w):
    """weighted mean over axis 1 that ignores NaN labels. Y [b,k,2], w [b,k] -> [b,2] (NaN if no valid label)"""
    ok = ~np.isnan(Y)
    ww = w[:, :, None] * ok
    den = ww.sum(1)
    num = (np.where(ok, Y, 0.0) * ww).sum(1)
    return np.where(den > 0, num / np.maximum(den, 1e-12), np.nan)


# ---------------------------------------- SUBSTRATE "AVERAGE" FEATURES ---
def te_stats(df):
    """per-substrate: row count, mean/std of log10_kcat and log10_km, plus ENZYME-BALANCED means
    (average over enzyme groups instead of rows, so enzymes with hundreds of mutants do not dominate)."""
    g = df.groupby("smiles").agg(n=("log10_kcat", "size"),
                                 kc=("log10_kcat", "mean"), km=("log10_km", "mean"),
                                 kc_sd=("log10_kcat", "std"), km_sd=("log10_km", "std"))
    if "_g" in df.columns:
        d = df.groupby(["smiles", "_g"]).agg(kc=("log10_kcat", "mean"), km=("log10_km", "mean")).reset_index()
        bal = d.groupby("smiles").agg(kcb=("kc", "mean"), kmb=("km", "mean"), ng=("kc", "size"))
        g = g.join(bal)
    else:
        g["kcb"], g["kmb"], g["ng"] = g["kc"], g["km"], g["n"]
    return g


def te_apply(smiles, stats):
    s = stats.reindex(smiles)
    return np.column_stack([s["n"].fillna(0).values, s["kc"].values, s["km"].values,
                            s["kc_sd"].values, s["km_sd"].values,
                            s["kcb"].values, s["kmb"].values, s["ng"].fillna(0).values]).astype(np.float32)


def sub_neighbor_feats(q_bits, ref_bits, ref_vals, k=5):
    """For each query substrate: how similar are the k most similar TRAINING substrates
    (Tanimoto on Morgan bits) and what are their weighted-mean kcat / km.
    This helps for the ~1/3 of validation substrates never seen in training."""
    nq = q_bits.shape[0]
    out = np.full((nq, 4), np.nan, dtype=np.float32)
    rs = ref_bits.sum(1)
    k = min(k, ref_bits.shape[0] - 1)
    for i in range(0, nq, 512):
        qb = q_bits[i:i + 512]
        inter = qb @ ref_bits.T
        union = qb.sum(1)[:, None] + rs[None, :] - inter
        sim = inter / np.maximum(union, 1)
        idx = np.argpartition(-sim, k, axis=1)[:, :k]
        s = np.take_along_axis(sim, idx, 1)
        w = s ** 2 + 1e-6
        vals = ref_vals[idx]
        out[i:i + 512, 0] = s.max(1)
        out[i:i + 512, 1:3] = _wmean(vals, w)
        out[i:i + 512, 3] = s.mean(1)
    return out


def te_full(smiles_values, smi_index, bits_u, stats, ref_bits):
    """9 substrate-level features per row: exact-substrate stats (5) + similar-substrate stats (4).
    smi_index: dict smiles -> row of bits_u; ref_bits aligned with stats.index."""
    exact = te_apply(smiles_values, stats)
    uq = list(dict.fromkeys(smiles_values))
    q = bits_u[[smi_index[s] for s in uq]]
    ref_vals = stats[["kc", "km"]].values.astype(np.float32)
    nb = sub_neighbor_feats(q, ref_bits, ref_vals)
    pos = {s: i for i, s in enumerate(uq)}
    nb_rows = nb[[pos[s] for s in smiles_values]]
    return np.hstack([exact, nb_rows]).astype(np.float32)


# ----------------------------------------------------------------- ASSEMBLE ---
def assemble(df, smi_unique, smi_mat, seq_unique, seq_mat, esm_mat=None, te=None):
    smi_idx = {s: i for i, s in enumerate(smi_unique)}
    seq_idx = {s: i for i, s in enumerate(seq_unique)}
    si = df["smiles"].map(smi_idx).values
    qi = df["sequence"].map(seq_idx).values
    parts = [smi_mat[si], seq_mat[qi]]
    if esm_mat is not None:
        parts.append(esm_mat[qi])
    if te is not None:
        parts.append(te)
    return np.hstack(parts).astype(np.float32)


# ------------------------------------- ESM NEIGHBOUR ("similar enzyme") FEATURES ---
def _knn_one(q_smiles, q_si, q_embn, ref_smiles, ref_si, ref_embn, ref_y, k=5, temp=5.0):
    """For each query row, look at training enzymes that ESM thinks are similar.
    cols 0-2: among training rows with the SAME substrate: best similarity, weighted mean kcat, weighted mean km
    cols 3-5: among ALL training enzymes (per-enzyme mean targets): best similarity, mean kcat, mean km
    q_si / ref_si = row index of each row's sequence inside q_embn / ref_embn (centered + length-normalised)."""
    from collections import defaultdict
    nq = len(q_smiles)
    out = np.full((nq, 6), np.nan, dtype=np.float32)
    q_si = np.asarray(q_si); ref_si = np.asarray(ref_si); ref_y = np.asarray(ref_y, dtype=np.float32)
    ref_by = defaultdict(list)
    for j, s in enumerate(ref_smiles):
        ref_by[s].append(j)
    q_by = defaultdict(list)
    for i, s in enumerate(q_smiles):
        q_by[s].append(i)
    for s, qi in q_by.items():
        rj = ref_by.get(s)
        if not rj:
            continue
        rj = np.array(rj); qi = np.array(qi)
        sim = q_embn[q_si[qi]] @ ref_embn[ref_si[rj]].T
        kk = min(k, len(rj))
        idx = np.argpartition(-sim, kk - 1, axis=1)[:, :kk]
        sv = np.take_along_axis(sim, idx, 1)
        w = np.exp(temp * sv)
        y = ref_y[rj][idx]
        out[qi, 0] = sv.max(1)
        out[qi, 1:3] = _wmean(y, w)
    # global: nearest enzymes, using each enzyme's mean targets (NaN labels ignored)
    n_ref = ref_embn.shape[0]
    ok = ~np.isnan(ref_y)
    ysum = np.zeros((n_ref, 2)); cnt = np.zeros((n_ref, 2))
    np.add.at(ysum, ref_si, np.where(ok, ref_y, 0.0)); np.add.at(cnt, ref_si, ok.astype(float))
    have = np.nonzero(cnt.sum(1))[0]
    ref_seq_y = np.where(cnt[have] > 0, ysum[have] / np.maximum(cnt[have], 1), np.nan).astype(np.float32)
    R = ref_embn[have]
    uq = np.unique(q_si)
    G = np.full((len(uq), 3), np.nan, dtype=np.float32)
    kk = min(k, len(have))
    for i in range(0, len(uq), 1024):
        sim = q_embn[uq[i:i + 1024]] @ R.T
        idx = np.argpartition(-sim, kk - 1, axis=1)[:, :kk]
        sv = np.take_along_axis(sim, idx, 1)
        w = np.exp(temp * sv)
        G[i:i + 1024, 0] = sv.max(1)
        G[i:i + 1024, 1:3] = _wmean(ref_seq_y[idx], w)
    pos = {int(s): i for i, s in enumerate(uq)}
    out[:, 3:6] = G[[pos[int(s)] for s in q_si]]
    return out


def normalize_emb(emb, mean):
    c = emb - mean
    return (c / np.maximum(np.linalg.norm(c, axis=1, keepdims=True), 1e-6)).astype(np.float32)


def knn_feats(q_smiles, q_si, q_embn, ref_smiles, ref_si, ref_embn, ref_y, ks=(5,), temp=5.0):
    """6 columns per value of k (see _knn_one)."""
    return np.hstack([_knn_one(q_smiles, q_si, q_embn, ref_smiles, ref_si, ref_embn, ref_y, k=k, temp=temp)
                      for k in ks])


# ---------------------------------------------------------------- VARIANTS ---
def fit_pca(emb, n):
    mean = emb.mean(0)
    _, _, vt = np.linalg.svd(emb - mean, full_matrices=False)
    return {"mean": mean.astype(np.float32), "comps": vt[:n].astype(np.float32)}


def apply_variant(X0, kn, p, smi_dim, seq_dim, esm_dim, ks_all=(5,), pca=None, n_blocks=1, er=None):
    """Shared by train.py and predict.py. X0 columns = [smiles | sequence | esm | substrate-stats]."""
    esm_lo = smi_dim + seq_dim
    drop = []
    if p.get("esm_blocks") is not None and esm_dim:
        bs = esm_dim // n_blocks
        for b in range(n_blocks):
            if b not in p["esm_blocks"]:
                drop += list(range(esm_lo + b * bs, esm_lo + (b + 1) * bs))
    if p.get("noesm") or p.get("esm_pca"):
        drop += list(range(esm_lo, esm_lo + esm_dim))
    if p.get("nodi"):
        drop += list(range(smi_dim + 22, smi_dim + seq_dim))      # the 400 dipeptide columns
    keep = np.setdiff1d(np.arange(X0.shape[1]), np.array(drop, dtype=int))
    parts = [X0[:, keep]]
    if p.get("esm_pca"):
        parts.append((X0[:, esm_lo:esm_lo + esm_dim] - pca["mean"]) @ pca["comps"].T)
    if p.get("knn"):
        cols = []
        for k in p.get("ks", (5,)):
            j = list(ks_all).index(k)
            cols.extend(range(6 * j, 6 * j + 6))
        parts.append(kn[:, cols])
    if p.get("res"):
        parts.append(er)
    return np.hstack(parts).astype(np.float32)


# ------------------------------------------------------------- MLP HELPERS ---
def mlp_prep_fit(X):
    X = X.astype(np.float32).copy()
    X[:, :N_MORGAN] = np.log1p(np.maximum(X[:, :N_MORGAN], 0))
    with np.errstate(all="ignore"):
        mean = np.nanmean(X, axis=0)
        std = np.nanstd(X, axis=0)
    mean = np.nan_to_num(mean); std = np.nan_to_num(std)
    std[std < 1e-6] = 1.0
    return {"mean": mean.astype(np.float32), "std": std.astype(np.float32)}


def mlp_prep(X, st):
    X = X.astype(np.float32).copy()
    X[:, :N_MORGAN] = np.log1p(np.maximum(X[:, :N_MORGAN], 0))
    Z = (X - st["mean"]) / st["std"]
    return np.clip(np.nan_to_num(Z, nan=0.0), -8, 8).astype(np.float32)


def mlp_forward(Z, layers):
    h = Z
    for i, (W, b) in enumerate(layers):
        h = h @ W + b
        if i < len(layers) - 1:
            h = np.maximum(h, 0)
    return h


# ------------------------------------------- ENZYME-EFFECT (residual) FEATURES ---
def enzyme_effect_ref(ref_smiles, ref_si, ref_y, n_ref, stats, m=3.0):
    """Split each label into  substrate effect + enzyme effect.  Returns, for every enzyme (row of the
    embedding matrix), its average residual = label - (shrunk) substrate mean.  NaN if the enzyme is absent."""
    ref_y = np.asarray(ref_y, dtype=np.float64)
    gm = np.nanmean(ref_y, axis=0)
    s = stats.reindex(ref_smiles)
    n = s["n"].values[:, None]
    mu = s[["kc", "km"]].values
    shr = (n * mu + m * gm) / (n + m)
    res = ref_y - shr
    ok = ~np.isnan(res)
    rs = np.zeros((n_ref, 2)); cnt = np.zeros((n_ref, 2))
    np.add.at(rs, ref_si, np.where(ok, res, 0.0)); np.add.at(cnt, ref_si, ok.astype(float))
    return np.where(cnt > 0, rs / np.maximum(cnt, 1), np.nan).astype(np.float32)


def enzyme_effect_feats(q_si, q_embn, ref_embn, ref_seq_res, k=10, temp=5.0):
    """How much faster / slower than average (for its substrate) are the training enzymes that ESM finds
    most similar to this one?  3 columns: best similarity, weighted mean residual kcat, residual km."""
    q_si = np.asarray(q_si)
    has = np.nonzero(~np.isnan(ref_seq_res).all(1))[0]
    R = ref_embn[has]; Y = ref_seq_res[has]
    uq = np.unique(q_si)
    G = np.full((len(uq), 3), np.nan, dtype=np.float32)
    kk = min(k, len(has))
    for i in range(0, len(uq), 1024):
        sim = q_embn[uq[i:i + 1024]] @ R.T
        idx = np.argpartition(-sim, kk - 1, axis=1)[:, :kk]
        sv = np.take_along_axis(sim, idx, 1)
        w = np.exp(temp * sv)
        G[i:i + 1024, 0] = sv.max(1)
        G[i:i + 1024, 1:3] = _wmean(Y[idx], w)
    pos = {int(s): i for i, s in enumerate(uq)}
    return G[[pos[int(s)] for s in q_si]]