"""Sparse k-mer index: TF-IDF-weighted, L2-normalised presence vectors for cosine search."""
import numpy as np
import scipy.sparse as sp

AA = 'ACDEFGHIKLMNPQRSTVWY'
LUT = np.full(256, 20, dtype=np.int64)
for i, c in enumerate(AA):
    LUT[ord(c)] = i
    LUT[ord(c.lower())] = i


def kmer_ids(seq, k):
    a = LUT[np.frombuffer(seq.encode('ascii', 'replace'), dtype=np.uint8)]
    if len(a) < k:
        return np.empty(0, dtype=np.int64)
    ids = np.zeros(len(a) - k + 1, dtype=np.int64)
    for j in range(k):
        ids = ids * 21 + a[j:len(a) - k + 1 + j]
    return np.unique(ids)


def presence_matrix(seqs, k):
    rows, cols = [], []
    for i, s in enumerate(seqs):
        ids = kmer_ids(s, k)
        rows.append(np.full(len(ids), i, dtype=np.int32))
        cols.append(ids)
    rows = np.concatenate(rows)
    cols = np.concatenate(cols)
    data = np.ones(len(rows), dtype=np.float32)
    return sp.csr_matrix((data, (rows, cols)), shape=(len(seqs), 21 ** k))


def weight(m, idf):
    m = m.multiply(idf[None, :]).tocsr().astype(np.float32)
    norms = np.sqrt(np.asarray(m.multiply(m).sum(axis=1)).ravel())
    norms[norms == 0] = 1
    return sp.diags(1 / norms).dot(m).tocsr().astype(np.float32)


class KmerIndex:
    def __init__(self, seqs, k=5):
        self.k = k
        m = presence_matrix(seqs, k)
        df = np.bincount(m.indices, minlength=m.shape[1]).astype(np.float32)
        self.idf = np.log((1 + m.shape[0]) / (1 + df)).astype(np.float32)
        self.T = weight(m, self.idf).T.tocsr()  # (features x n_train)

    def topk(self, seqs, k=50, chunk=256):
        """Return (indices, scores) arrays of shape (len(seqs), k)."""
        n = len(seqs)
        idx = np.full((n, k), -1, dtype=np.int64)
        sc = np.zeros((n, k), dtype=np.float32)
        for start in range(0, n, chunk):
            q = weight(presence_matrix(seqs[start:start + chunk], self.k), self.idf)
            s = (q @ self.T).tocsr()
            for r in range(s.shape[0]):
                lo, hi = s.indptr[r], s.indptr[r + 1]
                if hi == lo:
                    continue
                d, c = s.data[lo:hi], s.indices[lo:hi]
                m = min(k, hi - lo)
                part = np.argpartition(-d, m - 1)[:m]
                part = part[np.argsort(-d[part])]
                idx[start + r, :m] = c[part]
                sc[start + r, :m] = d[part]
        return idx, sc