"""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):
    """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
    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():
            h = model(**enc).last_hidden_state
        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
        e = (h * m.unsqueeze(-1)).sum(1) / m.sum(1, keepdim=True).clamp(min=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


# ---------------------------------------- SUBSTRATE "AVERAGE" FEATURES ---
def te_stats(df):
    """per-substrate: row count, mean and std of log10_kcat / log10_km in the training rows"""
    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"))
    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]).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]
        mean = (vals * w[:, :, None]).sum(1) / w.sum(1)[:, None]
        out[i:i + 512, 0] = s.max(1)
        out[i:i + 512, 1:3] = mean
        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)