"""Core random-indexing primitives used by the encoder."""

from __future__ import annotations

import hashlib
from dataclasses import dataclass

import numpy as np

from .types import FeatureLog, SparseVector


@dataclass(frozen=True)
class RIVectorSpec:
    """Deterministic RI basis-vector configuration."""

    dim: int = 2048
    nonzeros: int = 8
    seed: int = 17


def _stable_hash_u64(s: str, seed: int) -> int:
    """Hash a token into a stable unsigned 64-bit integer."""
    h = hashlib.blake2b((str(seed) + "::" + s).encode("utf-8"), digest_size=8).digest()
    return int.from_bytes(h, "little", signed=False)


def ri_vector_atomic(token: str, spec: RIVectorSpec) -> SparseVector:
    """Create the sparse RI basis vector for a single feature token."""
    if spec.dim <= 0:
        raise ValueError("dim must be > 0")
    if spec.nonzeros <= 0:
        return {}

    h = _stable_hash_u64(token, spec.seed)
    rng = np.random.default_rng(h)
    size = min(spec.nonzeros, spec.dim)
    idx = rng.choice(spec.dim, size=size, replace=False)
    signs = rng.choice([-1.0, 1.0], size=size, replace=True)
    return {int(i): float(s) for i, s in zip(idx, signs)}


def add_sparse(dest: SparseVector, src: SparseVector, weight: float = 1.0) -> None:
    """Accumulate a weighted sparse vector into another sparse vector."""
    for k, v in src.items():
        dest[k] = dest.get(k, 0.0) + weight * v


def sparse_to_dense(vec: SparseVector, dim: int) -> np.ndarray:
    """Convert a sparse dictionary vector into a dense float32 array."""
    x = np.zeros(dim, dtype=np.float32)
    for k, v in vec.items():
        if 0 <= k < dim:
            x[k] = np.float32(v)
    return x


def dense_l2_normalize(x: np.ndarray) -> np.ndarray:
    """L2-normalize a dense vector, preserving zero vectors as-is."""
    n = float(np.linalg.norm(x))
    if n <= 0:
        return x.astype(np.float32, copy=False)
    return (x / np.float32(n)).astype(np.float32, copy=False)


def char_ngrams(s: str, n_min: int = 3, n_max: int = 5) -> list[str]:
    """Return boundary-aware character n-grams for a token."""
    if n_min <= 0 or n_max < n_min:
        return []

    s2 = f"#{s}#"
    out: list[str] = []
    for n in range(n_min, n_max + 1):
        for i in range(0, max(0, len(s2) - n + 1)):
            out.append(s2[i : i + n])
    return out


def _sample_deterministic(items: list[str], k: int, seed_u64: int) -> list[str]:
    """Select up to ``k`` items using a stable hash-based ordering."""
    if k <= 0 or len(items) <= k:
        return items

    scored: list[tuple[int, str]] = []
    for it in items:
        h = _stable_hash_u64(it, int(seed_u64 & 0xFFFFFFFF))
        scored.append((h, it))
    scored.sort(key=lambda x: x[0])
    return [it for _, it in scored[:k]]


@dataclass(frozen=True)
class RIContextParams:
    """Feature extraction and weighting parameters for contextual RI encoding."""

    use_words: bool = True
    use_chargrams: bool = True
    gram_min: int = 3
    gram_max: int = 5
    max_grams_per_word: int = 40

    w_family: float = 2.0
    w_word: float = 1.0
    w_chargram: float = 0.3

    position_decay: float = 0.0


def ri_vector_context_sparse(
    family: str,
    text_norm: str,
    spec: RIVectorSpec,
    params: RIContextParams = RIContextParams(),
) -> tuple[SparseVector, FeatureLog]:
    """Encode normalized text into a sparse contextual RI vector.

    The sparse vector combines a family token, word features, and optional
    character n-grams using deterministic RI basis vectors.
    """
    vec: SparseVector = {}
    feats_log: FeatureLog = []

    fam_tok = f"fam::{family}"
    add_sparse(vec, ri_vector_atomic(fam_tok, spec), weight=params.w_family)
    feats_log.append((fam_tok, params.w_family))

    words = [w for w in text_norm.split() if w] if params.use_words else []
    if not words:
        raw_tok = f"raw::{family}::{text_norm}"
        add_sparse(vec, ri_vector_atomic(raw_tok, spec), weight=1.0)
        feats_log.append((raw_tok, 1.0))
        return vec, feats_log

    for i, w in enumerate(words):
        wpos = 1.0
        if params.position_decay > 0:
            wpos = 1.0 / (1.0 + params.position_decay * float(i))

        wt = params.w_word * wpos
        wtok = f"w::{family}::{w}"
        add_sparse(vec, ri_vector_atomic(wtok, spec), weight=wt)
        feats_log.append((wtok, wt))

        if params.use_chargrams and params.w_chargram != 0.0:
            grams = char_ngrams(w, params.gram_min, params.gram_max)
            grams = _sample_deterministic(
                grams,
                params.max_grams_per_word,
                seed_u64=_stable_hash_u64(w, spec.seed),
            )
            gwt = params.w_chargram * wpos
            for g in grams:
                gtok = f"g::{family}::{g}"
                add_sparse(vec, ri_vector_atomic(gtok, spec), weight=gwt)
                feats_log.append((gtok, gwt))

    return vec, feats_log


def ri_vector_context_dense(
    family: str,
    text_norm: str,
    spec: RIVectorSpec,
    params: RIContextParams = RIContextParams(),
) -> tuple[np.ndarray, FeatureLog]:
    """Encode normalized text into a dense L2-normalized RI vector."""
    vec_sparse, feats_log = ri_vector_context_sparse(family, text_norm, spec, params)
    dense = sparse_to_dense(vec_sparse, spec.dim)
    dense = dense_l2_normalize(dense)
    return dense, feats_log
