"""Feature extraction helpers used by the encoder and explanation layer."""

from __future__ import annotations

from .ri import RIContextParams, RIVectorSpec, _sample_deterministic, _stable_hash_u64, char_ngrams
from .types import FeatureMap


def tokenize_words(text_norm: str) -> list[str]:
    """Split normalized text into non-empty whitespace-delimited tokens."""
    return [w for w in text_norm.split() if w]


def extract_feature_weights(
    family: str,
    text_norm: str,
    params: RIContextParams,
    spec: RIVectorSpec,
) -> FeatureMap:
    """Build the weighted feature map for a normalized text input.

    The returned map mirrors the features that contribute to the RI vector:
    one family token, zero or more word tokens, and optional char n-grams.
    """
    out: FeatureMap = {}

    fam_tok = f"fam::{family}"
    out[fam_tok] = out.get(fam_tok, 0.0) + params.w_family

    words = tokenize_words(text_norm) if params.use_words else []
    if not words:
        raw_tok = f"raw::{family}::{text_norm}"
        out[raw_tok] = out.get(raw_tok, 0.0) + 1.0
        return out

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

        wtok = f"w::{family}::{w}"
        out[wtok] = out.get(wtok, 0.0) + params.w_word * wpos

        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 gram in grams:
                gtok = f"g::{family}::{gram}"
                out[gtok] = out.get(gtok, 0.0) + gwt

    return out
