"""Text primitives shared by every stage: normalization, token estimates, n-grams, MinHash, hashed features.

Pure standard library so the hosted preview needs no heavy installs. numpy is used when present to speed up MinHash.
"""
from __future__ import annotations

import hashlib
import math
import re
import struct
import unicodedata

try:  # optional speed-up for MinHash on large runs
    import numpy as _np
except Exception:  # pragma: no cover
    _np = None

WORD = re.compile(r"\w+", re.U)
_MERSENNE = (1 << 61) - 1
_MAXH = (1 << 32) - 1


def norm(text: str) -> str:
    """Lowercase, strip accents and punctuation, squeeze whitespace. Used for exact dedup and n-gram matching."""
    t = unicodedata.normalize("NFKD", text.lower())
    t = "".join(c for c in t if not unicodedata.combining(c))
    return " ".join(WORD.findall(t))


def words(text: str) -> list[str]:
    return WORD.findall(text.lower())


def est_tokens(text: str) -> int:
    """About four characters per token for English prose. Good enough to plan a budget, not to bill one."""
    return max(1, len(text) // 4)


def h64(s: str) -> int:
    return int.from_bytes(hashlib.blake2b(s.encode("utf-8", "ignore"), digest_size=8).digest(), "little")


def doc_hash(text: str) -> str:
    return hashlib.blake2b(norm(text).encode(), digest_size=12).hexdigest()


def ngrams(ws: list[str], n: int) -> list[str]:
    return [" ".join(ws[i:i + n]) for i in range(len(ws) - n + 1)] if len(ws) >= n else ([" ".join(ws)] if ws else [])


# ---------------------------------------------------------------- MinHash + LSH (near-duplicate detection)

class MinHasher:
    """MinHash over word 5-gram shingles. 112 permutations in 14 bands of 8 rows find pairs at roughly 0.8 Jaccard,
    the setting FineWeb used for near-duplicate removal."""

    def __init__(self, num_perm: int = 112, bands: int = 14, shingle: int = 5, seed: int = 1):
        assert num_perm % bands == 0
        self.num_perm, self.bands, self.rows, self.shingle = num_perm, bands, num_perm // bands, shingle
        import random
        rng = random.Random(seed)
        # a, b < 2**32 keep a*x + b inside 64 bits, so the numpy and pure-Python paths agree exactly
        self.a = [rng.randrange(1, _MAXH) for _ in range(num_perm)]
        self.b = [rng.randrange(0, _MAXH) for _ in range(num_perm)]
        if _np is not None:
            self._a = _np.array(self.a, dtype=_np.uint64)
            self._b = _np.array(self.b, dtype=_np.uint64)

    def signature(self, text: str) -> list[int]:
        sh = {h64(g) & _MAXH for g in ngrams(words(text), self.shingle)} or {0}
        if _np is not None:
            x = _np.fromiter(sh, dtype=_np.uint64)[:, None]
            v = ((x * self._a + self._b) % _np.uint64(_MERSENNE)) & _np.uint64(_MAXH)
            return v.min(axis=0).tolist()
        return [min(((a * x + b) % _MERSENNE) & _MAXH for x in sh) for a, b in zip(self.a, self.b)]

    def band_keys(self, sig: list[int]) -> list[bytes]:
        r = self.rows
        return [bytes([i]) + struct.pack(f"<{r}I", *sig[i * r:(i + 1) * r]) for i in range(self.bands)]


def jaccard_sig(a: list[int], b: list[int]) -> float:
    return sum(1 for x, y in zip(a, b) if x == y) / len(a)


class NearDupIndex:
    """In-memory LSH index. Holds one band table per band; fine for tens of millions of documents on a large box.
    Past that, shard by a hash prefix and run one index per shard."""

    def __init__(self, hasher: MinHasher | None = None, threshold: float = 0.8):
        self.h = hasher or MinHasher()
        self.threshold = threshold
        self.tables: list[dict] = [dict() for _ in range(self.h.bands)]
        self.sigs: dict[int, list[int]] = {}

    def check_add(self, key: int, text: str) -> int | None:
        """Return the key of an earlier near-duplicate, or None after adding this document."""
        return self.check_add_sig(key, self.h.signature(text))

    def check_add_sig(self, key: int, sig: list[int]) -> int | None:
        """Same, with a signature computed elsewhere (in a worker process)."""
        cands = set()
        for i, k in enumerate(self.h.band_keys(sig)):
            hit = self.tables[i].get(k)
            if hit is not None:
                cands.add(hit)
        for c in cands:
            if jaccard_sig(sig, self.sigs[c]) >= self.threshold:
                return c
        for i, k in enumerate(self.h.band_keys(sig)):
            self.tables[i].setdefault(k, key)
        self.sigs[key] = sig
        return None


# ---------------------------------------------------------------- hashed features (cheap classifier + relevance)

def features(text: str, dim: int = 1 << 18) -> dict[int, float]:
    """Hashed unigrams and bigrams, log-scaled and L2-normalized. Works for any language with word characters."""
    ws = words(text[:20000])
    counts: dict[int, float] = {}
    for g in ws + ngrams(ws, 2):
        k = h64(g) % dim
        counts[k] = counts.get(k, 0.0) + 1.0
    for k in counts:
        counts[k] = 1.0 + math.log(counts[k])
    n = math.sqrt(sum(v * v for v in counts.values())) or 1.0
    return {k: v / n for k, v in counts.items()}


def cosine(a: dict[int, float], b: dict[int, float]) -> float:
    if len(a) > len(b):
        a, b = b, a
    return sum(v * b.get(k, 0.0) for k, v in a.items())


def centroid(vecs: list[dict[int, float]]) -> dict[int, float]:
    c: dict[int, float] = {}
    for v in vecs:
        for k, x in v.items():
            c[k] = c.get(k, 0.0) + x
    n = math.sqrt(sum(x * x for x in c.values())) or 1.0
    return {k: x / n for k, x in c.items()}