"""Compositional HRR morpheme tokenizer.

Each word is represented as a single fixed-width "holistic" vector:

    word_vec(word) = bundle(
        prefix_role (x) bytes(prefix),
        root_role   (x) bytes(root),
        suffix_role (x) bytes(suffix),
    )

where ``(x)`` is circular convolution (``hrr.bind``) and ``bundle`` is normalized
superposition. Because the roles are unitary, ``unbind(word_vec, role)`` recovers
that role's filler — enabling algebraic morpheme manipulation (strip a prefix,
swap a suffix, build a vector for an out-of-vocabulary word from its pieces).

NOTE: this is an input **representation / embedding**, not a HuggingFace
text<->ID tokenizer. It emits dense vectors (one per word), intended to feed
HRR-native or experimental models.
"""
from __future__ import annotations

import re

import mlx.core as mx
import numpy as np

from .hrr import bind, bundle, make_unitary, normalize
from .morphemes import segment as _segment

_WORD_RE = re.compile(r"[A-Za-z]+|\S")


class MorphemeTokenizer:
    """Map text -> HRR morpheme vectors of width ``dim`` (deterministic given ``seed``)."""

    def __init__(self, dim: int = 2048, seed: int = 0):
        self.dim = dim
        self.seed = seed
        # Fixed random base vectors for each byte value; the "filler" alphabet.
        rng = np.random.default_rng(seed)
        byte_vectors = rng.normal(size=(256, dim)).astype(np.float32)
        self.byte_vectors = normalize(mx.array(byte_vectors))
        # Unitary role vectors for prefix / root / suffix slots.
        self.prefix_role = make_unitary(dim, seed=seed + 1_001)
        self.root_role = make_unitary(dim, seed=seed + 1_002)
        self.suffix_role = make_unitary(dim, seed=seed + 1_003)

    # -- morphemes ---------------------------------------------------------

    def segment(self, word: str) -> tuple[str, str, str]:
        """Return ``(prefix, root, suffix)`` for ``word``."""
        return _segment(word)

    def bytes_vector(self, text: str) -> mx.array:
        """Fixed vector for a string: normalized fold-bind of its byte vectors."""
        if not text:
            return mx.zeros((self.dim,), dtype=mx.float32)
        acc = self.byte_vectors[ord(text[0]) % 256]
        for ch in text[1:]:
            acc = bind(acc, self.byte_vectors[ord(ch) % 256])
        return normalize(acc)

    # -- composition -------------------------------------------------------

    def word_vector(self, word: str) -> mx.array:
        """Holistic HRR vector for ``word`` (composes its morphemes by role)."""
        prefix, root, suffix = _segment(word)
        pieces: list[mx.array] = []
        if prefix:
            pieces.append(bind(self.prefix_role, self.bytes_vector(prefix)))
        if root:
            pieces.append(bind(self.root_role, self.bytes_vector(root)))
        if suffix:
            pieces.append(bind(self.suffix_role, self.bytes_vector(suffix)))
        if not pieces:
            return self.bytes_vector(word)
        return bundle(*pieces)

    # -- text -> vectors ---------------------------------------------------

    def iter_vectors(self, text: str):
        """Yield one ``float16`` word vector per token in ``text`` (word or punct)."""
        for word in _WORD_RE.findall(text):
            vector = self.word_vector(word).astype(mx.float16)
            mx.eval(vector)
            yield mx.stop_gradient(vector)

    def encode(self, text: str) -> mx.array:
        """Stacked word vectors, shape ``(n_words, dim)`` in ``float16``."""
        vectors = list(self.iter_vectors(text))
        if not vectors:
            return mx.zeros((0, self.dim), dtype=mx.float16)
        return mx.stack(vectors)


# Backwards-compatible alias.
HolographicMorphemeTokenizer = MorphemeTokenizer