"""Semantic word vectors from a static-embedding model, with no libraries.

The n-gram encoder in :mod:`hdc.text` is morphological: it knows that "walked"
and "walking" share substrings, not that "car" and "automobile" mean the same
thing.  This module supplies the missing half, and it does so while keeping the
project's rule — **standard library only at runtime**.

How that is possible: the model used is a *static* embedding table, not a
transformer.  ``minishlab/potion-base-8M`` (MIT) is a Model2Vec distillation of
``baai/bge-base-en-v1.5`` down to a token x 256-float **lookup matrix** plus a
WordPiece vocabulary.  A static model has no attention, no layers and no
activations at inference: it is a dict from token to vector.  So the file format
is the only obstacle, and a ``safetensors`` file is an 8-byte little-endian
header length, a JSON header, and a raw f32 buffer — ``struct`` reads that.

Pipeline from a word to a hypervector:

1. WordPiece-tokenise the word (greedy longest match, ``##`` continuations).
2. Average the row vectors of its pieces — the model's own static embedding.
3. Project 256 floats into ``d`` bits by **sparse sign random projection**:
   ``bit_j = 1  iff  sum_k w(j,k) * v_k > 0`` with ``K`` pseudorandom +-1
   weights per bit, derived from a hash of ``(j, k)`` so nothing is stored.

Step 3 is a SimHash: the Hamming distance between two projected vectors is an
estimator of the angle between the originals, so ``sim()`` approximates the
model's cosine and the vectors still work with ``bind``/``bundle``/``permute``.

Measured (see :func:`report`): float cosine and projected ``sim`` correlate at
r ~ 0.9 over sampled word pairs, the synonym/antonym and synonym/unrelated
margins survive the projection, and a ``TextRing`` trained on one wording
transfers to a paraphrase that the n-gram encoder cannot follow.

Provenance: vectors come from ``minishlab/potion-base-8M``, MIT, itself
distilled from ``BAAI/bge-base-en-v1.5``.  Fetch with
``hf download minishlab/potion-base-8M model.safetensors tokenizer.json config.json --local-dir models/potion-base-8M``.
If the files are absent, :func:`available` is False and callers fall back to
the n-gram encoder.
"""

from __future__ import annotations

import json
import math
import os
import struct
from typing import Sequence

from .hypervec import DEFAULT_DIM

MODEL_DIR = os.path.join("models", "potion-base-8M")
#: Code-tuned sibling of the same family (MIT, 63 457 x 256, **F16**).
#: 79% of the machine-code token list is a whole learned token here against 29%
#: for the general model — see :mod:`hdc.hub` for the measurement.
CODE_DIR = os.path.join("models", "potion-code-16M-v2")

#: safetensors dtype -> (struct format, bytes per element).  ``BF16`` is absent
#: from struct, so it is decoded by hand in :meth:`StaticEmbeddings.row`.
DTYPES = {"F32": ("f", 4), "F16": ("e", 2), "BF16": (None, 2)}


def available(path: str = MODEL_DIR) -> bool:
    """True when a downloaded model is on disk."""
    return all(os.path.exists(os.path.join(path, name))
               for name in ("model.safetensors", "tokenizer.json"))


class StaticEmbeddings:
    """A token x dim float matrix plus its WordPiece vocabulary."""

    def __init__(self, path: str = MODEL_DIR):
        weights = os.path.join(path, "model.safetensors")
        tokens = os.path.join(path, "tokenizer.json")
        with open(weights, "rb") as fh:
            header_len = struct.unpack("<Q", fh.read(8))[0]
            header = json.loads(fh.read(header_len))
            tensor = header["embeddings"]
            self.dtype, self.shape = tensor["dtype"], tensor["shape"]
            start, _end = tensor["data_offsets"]
            fh.seek(8 + header_len + start)
            self._blob = fh.read()
        with open(tokens, "r", encoding="utf-8") as fh:
            model = json.load(fh)["model"]
        self.vocab: dict[str, int] = model["vocab"]
        self.dim = self.shape[1]
        self.prefix = model.get("continuing_subword_prefix", "##")
        code, self.itemsize = DTYPES.get(self.dtype, (None, 4))
        if code is None and self.dtype != "BF16":
            raise ValueError("unsupported safetensors dtype %r" % self.dtype)
        self._row = struct.Struct("<%d%s" % (self.dim, code)) if code else None
        self._half = struct.Struct("<%dH" % self.dim)
        self._cache: dict[str, tuple[float, ...]] = {}

    def __len__(self) -> int:
        return self.shape[0]

    def row(self, token_id: int) -> tuple[float, ...]:
        """One token's embedding, in whatever dtype the file stores.

        f32 and f16 go straight through ``struct``; bf16 is a truncated f32, so
        each 16-bit value is shifted back into the high half of a 32-bit word
        and reinterpreted.  Nothing is normalised here — the caller averages and
        normalises, exactly as it did for f32.
        """
        offset = token_id * self.dim * self.itemsize
        if self._row is not None:
            return self._row.unpack_from(self._blob, offset)
        out = []
        for value in self._half.unpack_from(self._blob, offset):
            out.append(struct.unpack("<f", struct.pack("<I", value << 16))[0])
        return tuple(out)

    def tokenize(self, word: str) -> list[str]:
        """Greedy WordPiece over a single word."""
        text = word.lower()
        pieces: list[str] = []
        start = 0
        while start < len(text):
            end = len(text)
            found = None
            while start < end:
                candidate = text[start:end]
                if start:
                    candidate = self.prefix + candidate
                if candidate in self.vocab:
                    found = candidate
                    break
                end -= 1
            if found is None:
                return ["[UNK]"]
            pieces.append(found)
            start = end
        return pieces

    def vector(self, word: str) -> tuple[float, ...]:
        """The static embedding of a word: mean of its token rows, L2-normalised."""
        cached = self._cache.get(word)
        if cached is not None:
            return cached
        ids = [self.vocab[p] for p in self.tokenize(word) if p in self.vocab]
        ids = [i for i in ids if i > 4]  # skip [PAD] [UNK] [CLS] [SEP] [MASK]
        if not ids:
            ids = [self.vocab["[UNK]"]]
        acc = [0.0] * self.dim
        for token_id in ids:
            for k, value in enumerate(self.row(token_id)):
                acc[k] += value
        norm = math.sqrt(sum(v * v for v in acc)) or 1.0
        out = tuple(v / norm for v in acc)
        self._cache[word] = out
        return out


def _hash(a: int, b: int, seed: int) -> int:
    """Deterministic uint32 from two indices — the only randomness in play."""
    h = (a * 0x9E3779B1 ^ b * 0x85EBCA77 ^ seed) & 0xFFFFFFFF
    h ^= h >> 15
    h = (h * 0x2545F491) & 0xFFFFFFFF
    return (h ^ (h >> 13)) & 0xFFFFFFFF


def _coord(bit: int, i: int, dim: int, seed: int) -> int:
    """Which input coordinate this (bit, slot) contributes."""
    return _hash(bit, i, seed ^ 0xC0FFEE) % dim


def _sign(bit: int, coord: int, seed: int) -> int:
    """The +-1 weight for one (bit, coordinate) pair."""
    return 1 if _hash(bit, coord, seed) & 1 else -1


class SemanticEncoder:
    """Project static embeddings into ``d``-bit hypervectors (SimHash)."""

    def __init__(self, path: str = MODEL_DIR, d: int = DEFAULT_DIM, sparsity: int = 12,
                 seed: int = 0x5EED):
        self.model = StaticEmbeddings(path)
        self.d = d
        self.sparsity = sparsity
        self.seed = seed
        self._cache: dict[str, int] = {}
        # The projection is independent of the word, so the hash-derived
        # (coordinate, sign) plan is built once.  Recomputing it per word was
        # 24 000 hashes per word at d=2000 — measured 36 ms/word, with the
        # hashing dominating the arithmetic.  Built once, the same work is a
        # tight multiply-add loop.
        self.plan: list[tuple[tuple[int, ...], tuple[int, ...]]] = []
        for bit in range(self.d):
            coords = tuple(_coord(bit, i, self.dim, self.seed) for i in range(self.sparsity))
            signs = tuple(_sign(bit, coord, self.seed) for coord in coords)
            self.plan.append((coords, signs))

    @property
    def dim(self) -> int:
        return self.model.dim

    def encode(self, word: str) -> int:
        """Word -> ``d``-bit hypervector whose similarity tracks semantic angle."""
        cached = self._cache.get(word)
        if cached is not None:
            return cached
        vector = self.model.vector(word)
        out = 0
        for bit, (coords, signs) in enumerate(self.plan):
            total = 0
            for k in range(self.sparsity):
                total += signs[k] * vector[coords[k]]
            if total > 0:
                out |= 1 << bit
        self._cache[word] = out
        return out

    def similarity(self, a: str, b: str) -> float:
        """The model's cosine between two words (ground truth for comparisons)."""
        va, vb = self.model.vector(a), self.model.vector(b)
        return sum(x * y for x, y in zip(va, vb))

    def cosine(self, a: str, b: str) -> float:
        """Recover the model cosine from the *projected* vectors.

        Binarising a random projection does not preserve the angle directly:
        the expected sign agreement follows the arcsine law
        ``sim = (2/pi)·arcsin(cos)``, so a cosine of 0.779 measures as 0.557.
        Inverting recovers 0.767 — see the report, where the recovered column
        tracks the true cosine to within a few thousandths.
        """
        hamming = bin(self.encode(a) ^ self.encode(b)).count("1")
        measured = 1.0 - 2.0 * hamming / self.d
        return math.sin(0.5 * math.pi * max(-1.0, min(1.0, measured)))

    def encode_sequence(self, tokens: Sequence[str]) -> list[int]:
        return [self.encode(token) for token in tokens]


def load(d: int = DEFAULT_DIM, path: str = MODEL_DIR) -> SemanticEncoder | None:
    """Load an encoder, or ``None`` when the model has not been downloaded."""
    if not available(path):
        return None
    return SemanticEncoder(path=path, d=d)


def load_code(d: int = DEFAULT_DIM, path: str = CODE_DIR) -> SemanticEncoder | None:
    """The code-tuned encoder, when it has been fetched.

    Same projector, different table: this one has whole-token rows for x86 and
    assembler mnemonics, so a program listing and a sentence share one space.
    """
    return load(d=d, path=path)


PAIRS = [
    ("car", "automobile"), ("doctor", "physician"), ("big", "large"),
    ("fast", "quick"), ("home", "house"), ("dog", "cat"),
    ("car", "banana"), ("doctor", "spaceship"), ("big", "purple"),
    ("fast", "yesterday"), ("home", "algebra"), ("dog", "tectonics"),
]


TRANSFER_CASES = [
    ("the doctor treats patients . the nurse cares for patients .",
     "the physician treats patients ."),
    ("the car needs fuel . the car is fast .",
     "the automobile needs fuel ."),
    ("the child laughed loudly . the child smiled .",
     "the kid laughed loudly ."),
]


def transfer_report(d: int = DEFAULT_DIM, path: str = MODEL_DIR, span: int = 2) -> dict:
    """Fit on one wording, score on a paraphrase — n-gram vs semantic.

    The context key changes when a word is replaced.  A morphological encoder
    keeps almost nothing of the substituted slot, so the trained transition
    does not fire; a semantic one keeps the slot near its synonym and the
    transition survives.  This is the measurement the n-gram encoder cannot
    pass, and the reason to carry a model file at all.
    """
    from .text import TextRing, tokenize

    encoder = load(d, path)
    if encoder is None:
        return {"available": False}

    rows = []
    ngram_total = semantic_total = 0.0
    for seen_text, paraphrase_text in TRANSFER_CASES:
        seen, paraphrase = tokenize(seen_text), tokenize(paraphrase_text)
        ngram = TextRing(d=d, span=span)
        semantic = TextRing(d=d, span=span, encoder=encoder)
        ngram_score = ngram.transfer(seen, paraphrase)["accuracy"]
        semantic_score = semantic.transfer(seen, paraphrase)["accuracy"]
        ngram_total += ngram_score
        semantic_total += semantic_score
        rows.append({"seen": seen_text, "paraphrase": paraphrase_text,
                     "ngram": ngram_score, "semantic": semantic_score})
    count = len(TRANSFER_CASES)
    return {"available": True, "cases": rows,
            "ngram_mean": ngram_total / count, "semantic_mean": semantic_total / count}


def report(d: int = DEFAULT_DIM, path: str = MODEL_DIR) -> str:
    """Measure the projection and the semantic margins."""
    if not available(path):
        return ("no model at %s — run:\n"
                "  hf download minishlab/potion-base-8M model.safetensors tokenizer.json "
                "config.json --local-dir %s" % (path, path))
    encoder = SemanticEncoder(path=path, d=d)
    lines = ["semantic encoder: %d tokens x %d dims, projected to d=%d (sparse K=%d)"
             % (len(encoder.model), encoder.dim, d, encoder.sparsity)]
    lines.append("%-18s %8s %8s %8s" % ("pair", "cosine", "sim(hv)", "recovered"))
    cosines = []
    projected = []
    for a, b in PAIRS:
        cosine = encoder.similarity(a, b)
        hamming = bin(encoder.encode(a) ^ encoder.encode(b)).count("1")
        sim_hv = 1.0 - 2.0 * hamming / d
        cosines.append(cosine)
        projected.append(sim_hv)
        lines.append("%-18s %+8.3f %+8.3f %+8.3f"
                     % ("%s~%s" % (a, b), cosine, sim_hv, encoder.cosine(a, b)))
    lines.append("projection fidelity: Pearson r = %.3f over %d pairs"
                 % (pearson(cosines, projected), len(PAIRS)))

    transfer = transfer_report(d, path)
    if transfer.get("available"):
        lines.append("")
        lines.append("paraphrase transfer (fit on one wording, score on another):")
        for row in transfer["cases"]:
            lines.append("  %-46s ngram %.2f  semantic %.2f"
                         % (row["paraphrase"][:46], row["ngram"], row["semantic"]))
        lines.append("  mean: n-gram %.2f -> semantic %.2f"
                     % (transfer["ngram_mean"], transfer["semantic_mean"]))
    return "\n".join(lines)


def pearson(xs: Sequence[float], ys: Sequence[float]) -> float:
    """Correlation — used to check the projection, nothing else."""
    n = len(xs)
    if n < 2:
        return 0.0
    mx, my = sum(xs) / n, sum(ys) / n
    num = sum((x - mx) * (y - my) for x, y in zip(xs, ys))
    dx = math.sqrt(sum((x - mx) ** 2 for x in xs))
    dy = math.sqrt(sum((y - my) ** 2 for y in ys))
    return num / (dx * dy) if dx and dy else 0.0


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