"""Small deterministic utilities used across GRM."""

from __future__ import annotations

import hashlib
import re
from collections.abc import Iterable

import torch
from torch import Tensor


def stable_hash(value: str, modulo: int) -> int:
    if modulo <= 0:
        raise ValueError("modulo must be positive")
    digest = hashlib.blake2b(value.encode("utf-8"), digest_size=8).digest()
    return int.from_bytes(digest, "big") % modulo


def tokenize(text: str | None) -> tuple[str, ...]:
    if not text:
        return ("<empty>",)
    return tuple(re.findall(r"[A-Za-z0-9_]+|[^\w\s]", text.lower()))


def strings_to_ids(values: Iterable[str], size: int) -> Tensor:
    return torch.tensor(
        [stable_hash(value, size) for value in values], dtype=torch.long
    )


def soft_cap(logits: Tensor, cap: float) -> Tensor:
    """``cap * tanh(logits / cap)`` in float32; the identity when ``cap <= 0``.

    Bounds attention scores smoothly.  Full-epoch runs blew up at the first
    attention layer each time: query/key weights grew until the scores
    saturated, and the pre-clip gradient norm went from single digits to
    thousands.  Near-identity for scores well below the cap, so a checkpoint
    trained without it behaves almost the same with it.
    """
    if cap <= 0:
        return logits
    return cap * torch.tanh(logits.float() / cap)


def qk_normalize(heads: Tensor) -> Tensor:
    """Scale each head vector to norm ``sqrt(head_dim)`` (QK normalisation).

    Applied to queries and keys, the usual ``q.k / sqrt(d)`` score becomes
    ``sqrt(d) * cosine``: bounded however large the projection weights grow.
    The soft cap bounded the scores but not that growth, and the first PAT
    layer's query/key gradients still climbed (median 13 -> ~360 between 190K
    and 245K steps).
    """
    dim = heads.shape[-1]
    unit = torch.nn.functional.normalize(heads.float(), dim=-1)
    return (unit * dim**0.5).to(heads.dtype)