"""Batched, vectorised byte n-gram hashing (numpy) - the shared input of the tier-1 HTTP model at training time
(torch EmbeddingBag) and inference time (numpy). Same function both sides, so the buckets always agree.

A batch of texts becomes one uint8 array with a 0 byte between texts; for each n-gram order the rolling polynomial
hash of every window is computed with a few whole-array operations, windows that cross a text boundary are dropped,
and a 64-bit mix maps the hash to a bucket. No per-character python, so cost is dominated by memory bandwidth.
"""
from __future__ import annotations

import numpy as np

P = np.uint64(0x100000001B3)            # FNV-64 prime
SEED = np.uint64(0x9E3779B97F4A7C15)
M1 = np.uint64(0xBF58476D1CE4E5B9)       # splitmix64 finaliser
M2 = np.uint64(0x94D049BB133111EB)
ORDERS = (2, 3, 4, 5, 6)
with np.errstate(over="ignore"):
    ORDER_SEEDS = {n: np.uint64(SEED * np.uint64(n)) for n in range(1, 17)}   # wraps mod 2^64 by design


def _mix(h: np.ndarray) -> np.ndarray:
    h = h ^ (h >> np.uint64(30))
    h = h * M1
    h = h ^ (h >> np.uint64(27))
    h = h * M2
    return h ^ (h >> np.uint64(31))


def pack(texts: list, max_bytes: int = 1024) -> tuple[np.ndarray, np.ndarray]:
    """texts -> (bytes with a 0 separator after each text, start offset of each text). Texts are cut at max_bytes."""
    enc = [(t if isinstance(t, (bytes, bytearray)) else str(t).encode("utf-8", "replace"))[:max_bytes].replace(b"\0", b" ")
           for t in texts]
    buf = np.frombuffer(b"\0".join(enc) + b"\0", dtype=np.uint8)
    lens = np.fromiter((len(e) for e in enc), dtype=np.int64, count=len(enc))
    starts = np.concatenate(([0], np.cumsum(lens + 1)[:-1])) if len(enc) else np.zeros(0, np.int64)
    return buf, starts


@np.errstate(over="ignore")
def ngram_buckets(texts: list, buckets: int, orders=ORDERS, max_bytes: int = 1024) -> tuple[np.ndarray, np.ndarray]:
    """-> (bucket index per n-gram, text index per n-gram), both int64, grouped by order then position.
    Texts shorter than the smallest order contribute nothing (callers treat an empty bag as all-zero)."""
    buf, starts = pack(texts, max_bytes)
    n_all = buf.size
    b64 = buf.astype(np.uint64)
    sep = (buf == 0).astype(np.int32)
    sep_cum = np.concatenate(([0], np.cumsum(sep)))                 # separators before position i
    doc_of = np.searchsorted(starts, np.arange(n_all), side="right") - 1
    h = b64 * P
    out_idx, out_doc = [], []
    for n in range(1, max(orders) + 1):
        if n > 1:
            h = h[:-1] * P + b64[n - 1:]                             # hash of buf[i:i+n]
        if n not in orders:
            continue
        m = n_all - n + 1
        ok = (sep_cum[n:n + m] - sep_cum[:m]) == 0                   # no separator inside the window
        hv = _mix(h[:m][ok] ^ ORDER_SEEDS[n])
        out_idx.append((hv % np.uint64(buckets)).astype(np.int64))
        out_doc.append(doc_of[:m][ok])
    if not out_idx:
        return np.zeros(0, np.int64), np.zeros(0, np.int64)
    return np.concatenate(out_idx), np.concatenate(out_doc)


def bag_mean(emb: np.ndarray, idx: np.ndarray, doc: np.ndarray, n_docs: int, scale: float | np.ndarray = 1.0) -> np.ndarray:
    """Mean of embedding rows per text (EmbeddingBag mode='mean'); emb may be int8 with a per-column `scale`."""
    d = emb.shape[1]
    out = np.zeros((n_docs, d), np.float32)
    if idx.size:
        # ngram_buckets output is sorted by text inside each order block: segment sums with reduceat (C speed),
        # instead of np.add.at (~5x slower)
        cuts = np.concatenate(([0], np.flatnonzero(np.diff(doc) < 0) + 1, [doc.size]))
        for s, e in zip(cuts[:-1], cuts[1:]):
            db = doc[s:e]
            first = np.flatnonzero(np.concatenate(([True], db[1:] != db[:-1])))
            out[db[first]] += np.add.reduceat(emb[idx[s:e]], first, axis=0, dtype=np.float32)
        cnt = np.bincount(doc, minlength=n_docs).astype(np.float32)
        out /= np.maximum(cnt, 1)[:, None]
    return out * np.asarray(scale, np.float32)


def offsets_for_torch(doc: np.ndarray, n_docs: int) -> tuple[np.ndarray, np.ndarray]:
    """(order of n-grams sorted by text, offsets) for torch.nn.EmbeddingBag(input, offsets)."""
    order = np.argsort(doc, kind="stable")
    counts = np.bincount(doc, minlength=n_docs)
    return order, np.concatenate(([0], np.cumsum(counts)[:-1]))