"""Text encoding for the ring: words as composed hypervectors.

The ring controller needs states.  Handing it ``hvc(word)`` works and is exact,
but it makes every word unrelated to every other word, so a context that has
never been seen predicts nothing.  Composing a word out of its character
n-grams fixes that: "walked" and "walking" share n-grams, so their vectors are
similar before any training has happened, and cleanup can fall back to the
nearest known form instead of failing.

    word  = bundle([bind(POS_i, atom(ngram_i)) for each char n-gram])

Position binding matters: without it "listen" and "silent" would be identical,
since they have the same n-gram multiset.  With it, they are not.

A context is then a record of recent words in role slots, which is exactly what
:class:`~hdc.state.PredictiveModel` wants as its key.

Measured on the corpus in :func:`evaluate` at d=2000: order-1 contexts give
67% top-1 next-token accuracy, order-2 gives 83% (the rest is genuine
two-way ambiguity in the corpus, not crosstalk), and composition keeps
inflections close — ``walked``~``walking`` scores +0.320 against +0.085 for
``walked``~``banana``.  That similarity is what lets a text-bound ring fall
back to a nearest known form instead of failing outright on a novel word.

Honest limits: this is a character-n-gram encoder, not a semantic model.  Two
synonyms are no closer than two random words.  The vector for a word says
"these substrings, in these positions", nothing more.
"""

from __future__ import annotations

import re
from typing import Iterable, Sequence

from .hypervec import DEFAULT_DIM, bind, bundle, hvc, permute, sim
from .memory_modes import HolographicMemory
from .state import PredictiveModel

_TOKEN = re.compile(r"[a-z']+|[.,!?;:]")

_START = "<"
_END = ">"


def tokenize(text: str) -> list[str]:
    """Lowercase word/punctuation tokens."""
    return _TOKEN.findall(text.lower())


def char_ngrams(word: str, order: int = 3) -> list[str]:
    """Character n-grams with boundary markers, so prefixes and suffixes count."""
    padded = _START * (order - 1) + word + _END * (order - 1)
    return [padded[i:i + order] for i in range(len(padded) - order + 1)]


def encode_word(word: str, d: int = DEFAULT_DIM, order: int = 3) -> int:
    """A word as a bundle of position-bound character n-grams."""
    parts = []
    for i, gram in enumerate(char_ngrams(word, order)):
        parts.append(bind(permute(hvc("text/pos", d), i, d), hvc("text/gram/" + gram, d), d))
    return bundle(parts, d) if parts else hvc("text/empty", d)


def encode_context(tokens: Sequence[str], d: int = DEFAULT_DIM, order: int = 3,
                   span: int = 2, encoder=None) -> int:
    """Recent tokens as a role-bound context record.

    ``span`` words of history, most recent first, each bound to its position —
    so order-1 and order-2 contexts are different vectors, never a superset.
    Pass an :class:`~hdc.semantic.SemanticEncoder` as ``encoder`` to make the
    context vectors semantic rather than morphological.
    """
    recent = list(tokens)[-span:]
    parts = []
    for back, token in enumerate(reversed(recent)):
        role = hvc("text/context/%d" % back, d)
        word = encoder.encode(token) if encoder is not None else encode_word(token, d, order)
        parts.append(bind(role, word, d))
    return bundle(parts, d) if parts else 0


def encode_text(text: str, d: int = DEFAULT_DIM, order: int = 3) -> list[tuple[str, int]]:
    """``(token, vector)`` pairs for a whole text."""
    return [(token, encode_word(token, d, order)) for token in tokenize(text)]


class TextRing:
    """A predictive model over word contexts, for use inside the ring."""

    def __init__(self, d: int = DEFAULT_DIM, order: int = 3, span: int = 2, encoder=None):
        self.d = d
        self.order = order
        self.span = span
        self.encoder = encoder
        self.model = PredictiveModel(d=d)
        self.vocabulary: dict[str, int] = {}
        self.contexts = 0

    def _word(self, token: str) -> int:
        vector = self.vocabulary.get(token)
        if vector is None:
            vector = (self.encoder.encode(token) if self.encoder is not None
                      else encode_word(token, self.d, self.order))
            self.vocabulary[token] = vector
            self.model.add_symbol(vector)
        return vector

    def _context(self, tokens: Sequence[str]) -> int:
        return encode_context(tokens, self.d, self.order, self.span, self.encoder)

    def fit(self, tokens: Sequence[str]) -> None:
        """Learn every context -> next-token transition in the token list."""
        for i in range(self.span, len(tokens) - 1):
            key = self._context(tokens[i - self.span + 1:i + 1])
            self.model.learn(key, "word", self._word(tokens[i + 1]))
            self.contexts += 1

    def predict(self, tokens: Sequence[str]) -> tuple[str | None, float]:
        """Most likely next token for the given history."""
        key = self._context(tokens)
        predicted, score = self.model.predict(key, "word")
        if predicted is None:
            return None, 0.0
        for token, vector in self.vocabulary.items():
            if vector == predicted:
                return token, score
        return None, score

    def accuracy(self, tokens: Sequence[str]) -> float:
        """Top-1 next-token accuracy over a sequence."""
        hits = total = 0
        for i in range(self.span, len(tokens) - 1):
            got, _ = self.predict(tokens[i - self.span + 1:i + 1])
            hits += got == tokens[i + 1]
            total += 1
        return hits / total if total else 0.0

    def transfer(self, seen: Sequence[str], paraphrased: Sequence[str]) -> dict:
        """Fit on one wording, score on another.

        This is the measurement that separates a semantic encoder from a
        morphological one: the context key changes when a word is replaced, and
        only a semantic encoder keeps it near the trained one.
        """
        self.fit(seen)
        return {
            "accuracy": self.accuracy(paraphrased),
            "seen": " ".join(seen),
            "paraphrased": " ".join(paraphrased),
        }


CORPUS = (
    "the cat sat on the mat . the cat ate the fish . "
    "the dog sat on the mat . the dog ate the bone . "
    "the cat sat on the mat"
)


def evaluate(d: int = DEFAULT_DIM, order: int = 3) -> dict:
    """Measure order-1 vs order-2 contexts, and whether composition generalises."""
    tokens = tokenize(CORPUS)
    results = {}
    for span in (1, 2):
        ring = TextRing(d=d, order=order, span=span)
        ring.fit(tokens)
        results["order%d_accuracy" % span] = ring.accuracy(tokens)
        results["order%d_vocab" % span] = len(ring.vocabulary)
        results["order%d_contexts" % span] = ring.contexts

    walked = encode_word("walked", d, order)
    walking = encode_word("walking", d, order)
    unrelated = encode_word("banana", d, order)
    results["similarity_inflections"] = sim(walked, walking, d)
    results["similarity_unrelated"] = sim(walked, unrelated, d)
    return results


def report(d: int = DEFAULT_DIM) -> str:
    """Printable encoder tour."""
    stats = evaluate(d)
    lines = ["text encoder, d=%d" % d]
    lines.append("order-1 contexts: %.0f%% top-1, %d tokens"
                 % (100 * stats["order1_accuracy"], stats["order1_vocab"]))
    lines.append("order-2 contexts: %.0f%% top-1, %d transitions"
                 % (100 * stats["order2_accuracy"], stats["order2_contexts"]))
    lines.append("composition keeps inflections close: walked~walking %+.3f, walked~banana %+.3f"
                 % (stats["similarity_inflections"], stats["similarity_unrelated"]))
    ring = TextRing(d=d, span=2)
    ring.fit(tokenize(CORPUS))
    history = tokenize("the cat sat on the")
    token, score = ring.predict(history)
    lines.append("predict after %r: %r (cleanup %.2f)" % (" ".join(history), token, score))
    return "\n".join(lines)


if __name__ == "__main__":  # pragma: no cover - manual inspection
    print(report())