"""Reading a direction back as words: the inverse of the transport.

`Lens` maps words to a direction. `Vocabulary` maps a direction back to the words whose own
transported directions align with it, which is how an axis gets a name — and how you find out
that it has none.

    voc = ulens.Lens("qwen3.5-4b", block=23).vocabulary()
    voc.top(axis, k=12)                  # the words this direction reads as
    print(voc.report(cone))              # every axis, both ends

Cheaply. Ranking by cosine against every transported row looks like it needs the full
transported vocabulary `D = unit((W_U * gamma) @ J)`, which is gigabytes, but it factorises:

    numerator_v   = (W_U[v] * gamma) . (J @ u)          one matvec, streamed
    denominator_v = || (W_U[v] * gamma) @ J ||          independent of u

so only the denominators have to be precomputed, and they are one float per token:
`sqrt(r_v @ (J J^T) @ r_v)` with `r_v = W_U[v] * gamma`, computed in chunks and cached at
about a megabyte instead of gigabytes. After that each decode is one streamed pass over the
unembedding. `cosine=False` skips the denominators altogether.

Two caveats that the numbers make unavoidable:

* **Transported rows share a large common component**, so raw cosines return much the same
  high-frequency words for unrelated directions. `purifier()` estimates the shared direction
  and its leading principal components, and `top(..., purify=...)` removes them.
* **An axis decoding to nothing is a real outcome.** Compare any top-k against
  `sample_directions()` — what a random direction of the same vocabulary decodes to — before
  reading meaning into it.
"""
from __future__ import annotations

import json
import re
from dataclasses import dataclass
from pathlib import Path

import numpy as np

from ._math import unit
from .models import ModelSpec
from .models import get as _model

#: words a decode will report: alphabetic, at least this long
ALPHA = re.compile(r"^[A-Za-z]+$")
MIN_LEN = 3

#: rows per streamed chunk. The denominator pass is O(chunk * hidden^2) and the decode pass
#: O(chunk * hidden), so the former runs in smaller chunks to bound peak memory.
NORM_CHUNK = 8192
DECODE_CHUNK = 65536


@dataclass(frozen=True)
class Hit:
    """One decoded token."""
    word: str
    score: float
    token_id: int

    def __repr__(self) -> str:
        return f"{self.word}({self.score:+.3f})"


@dataclass(frozen=True)
class AxisDecode:
    """Both ends of one axis, as words."""
    name: str
    positive: list[Hit]
    negative: list[Hit]

    def words(self, sign: str = "+") -> list[str]:
        return [h.word for h in (self.positive if sign == "+" else self.negative)]

    def __repr__(self) -> str:
        return f"AxisDecode({self.name!r}, +{self.words('+')[:4]}, -{self.words('-')[:4]})"


def common_component(D: np.ndarray, k: int = 3):
    """`(mu, V[k, hidden])`: the shared direction of unit rows `D`, and the leading principal
    components around it. Computed through the `hidden x hidden` covariance, so the number of
    rows only enters linearly."""
    X = np.asarray(D, np.float32)
    mu = X.mean(0)
    if k <= 0:
        return unit(mu), np.zeros((0, X.shape[1]), np.float32)
    Xc = X - mu
    C = (Xc.T @ Xc) / max(len(Xc) - 1, 1)
    _, V = np.linalg.eigh(C)
    return unit(mu), np.ascontiguousarray(V[:, ::-1][:, :k].T, np.float32)


def purify(v: np.ndarray, mu: np.ndarray, V: np.ndarray) -> np.ndarray:
    """Remove the shared direction and the leading principal components from `v`.

    Ranking a purified direction against the raw vocabulary is the same ordering as ranking
    the raw direction against a purified vocabulary, so this is applied to the direction.
    """
    v = np.asarray(v, np.float32)
    mu = np.asarray(mu, np.float32)
    v = v - (v @ mu) * mu
    V = np.asarray(V, np.float32)
    return v - V.T @ (V @ v) if V.size else v


class Vocabulary:
    """Every token of a model, as a transported direction, without materialising them."""

    def __init__(self, model: str | ModelSpec, block: int | None = None,
                 cache_dir: str | Path | None = None):
        self.spec = _model(model)
        self.block = self.spec.readout_block if block is None else int(block)
        self._cache_dir = Path(cache_dir) if cache_dir else None
        self._J: np.ndarray | None = None
        self._gamma: np.ndarray | None = None
        self._norms: np.ndarray | None = None
        self._n: int | None = None
        self._tok = None
        self._word: dict[int, str] = {}

    def __repr__(self) -> str:
        return f"Vocabulary({self.spec.key!r}, block={self.block})"

    # ------------------------------------------------------------------ resources
    @property
    def J(self) -> np.ndarray:
        if self._J is None:
            from . import lensfile as _lensfile
            self._J = _lensfile.matrix(self.spec.lens, self.block)
        return self._J

    @property
    def n_tokens(self) -> int:
        if self._n is None:
            from . import unembed
            self._n = unembed.unembedding_shape(self.spec.hf, self.spec.revision)[0]
        return self._n

    def _open(self):
        """`(handle, weight key, gamma)` for the unembedding, read lazily, never copied."""
        from . import unembed
        from ._deps import require
        safe_open = require("safetensors").safe_open
        names, shard = unembed._keys_and_shard(self.spec.hf, self.spec.revision)
        w_key = (unembed._pick(names, "lm_head.weight")
                 or unembed._pick(names, "embed_tokens.weight"))
        if self._gamma is None:
            g_key = unembed._pick(names, "model.norm.weight", ".norm.weight")
            with safe_open(shard(g_key), framework="pt") as f:
                self._gamma = f.get_tensor(g_key).float().numpy().astype(np.float32)
        return safe_open(shard(w_key), framework="pt"), w_key, self._gamma

    def _rows(self, chunk: int):
        """Yield `(start, (W_U[start:stop] * gamma))` as float32 `[c, hidden]`."""
        handle, w_key, gamma = self._open()
        with handle as f:
            sl = f.get_slice(w_key)
            n = int(sl.get_shape()[0])
            for i in range(0, n, chunk):
                R = sl[i:min(i + chunk, n), :].float().numpy().astype(np.float32)
                yield i, R * gamma[None, :]

    @property
    def cache_path(self) -> Path:
        from .basis import cache_dir
        root = self._cache_dir or cache_dir()
        return Path(root) / "vocab_norms" / f"{self.spec.key}-b{self.block}.npy"

    @property
    def norms(self) -> np.ndarray:
        """`[V]` row norms of the transported vocabulary. Computed once, then cached.

        `||r_v @ J||^2 = r_v @ (J J^T) @ r_v`, so the `V x hidden` product is never formed.
        """
        if self._norms is not None:
            return self._norms
        p = self.cache_path
        if p.exists():
            self._norms = np.load(p)
            return self._norms
        M = (self.J @ self.J.T).astype(np.float32)
        parts = [np.sqrt(np.maximum(((R @ M) * R).sum(1), 0.0))
                 for _, R in self._rows(NORM_CHUNK)]
        self._norms = np.concatenate(parts).astype(np.float32)
        p.parent.mkdir(parents=True, exist_ok=True)
        np.save(p, self._norms)
        return self._norms

    def word(self, token_id: int) -> str:
        """The decoded, stripped string of one token. Decoded lazily and cached, because a
        decode only ever inspects the few hundred candidates it scans."""
        t = int(token_id)
        if t not in self._word:
            if self._tok is None:
                from ._deps import require
                AutoTokenizer = require("transformers").AutoTokenizer
                kw = {"revision": self.spec.revision} if self.spec.revision else {}
                self._tok = AutoTokenizer.from_pretrained(self.spec.hf, **kw)
            self._word[t] = self._tok.decode([t]).strip()
        return self._word[t]

    # ------------------------------------------------------------------ decoding
    def scores(self, direction, cosine: bool = True, purify_with=None) -> np.ndarray:
        """Alignment of every token's transported direction with `direction`, `[V]`.

        `purify_with` is a `(mu, V)` pair from `purifier()`; it is applied to the direction.
        """
        u = np.asarray(direction, np.float32).reshape(-1)
        if u.shape[0] != self.spec.hidden_size:
            raise ValueError(f"direction is {u.shape[0]}-wide, "
                             f"{self.spec.key} is {self.spec.hidden_size}")
        if purify_with is not None:
            u = purify(u, *purify_with)
        v = (self.J @ u).astype(np.float32)
        s = np.concatenate([R @ v for _, R in self._rows(DECODE_CHUNK)]).astype(np.float32)
        if cosine:
            s = s / (self.norms + 1e-9) / (np.linalg.norm(u) + 1e-9)
        return s

    def top(self, direction, k: int = 12, sign: str = "+", min_len: int = MIN_LEN,
            alpha_only: bool = True, scan: int = 40, cosine: bool = True,
            purify_with=None) -> list[Hit]:
        """The `k` tokens best aligned with `direction` (`sign="-"` for the other end).

        Candidates are scanned `k * scan` deep and filtered to alphabetic words of at least
        `min_len` characters, deduplicated case-insensitively, which is the convention the
        reported decodes use: a raw top-k is mostly fragments and punctuation.
        """
        s = self.scores(direction, cosine=cosine, purify_with=purify_with)
        order = np.argsort(-s if sign == "+" else s)
        hits: list[Hit] = []
        seen: set[str] = set()
        for t in order[: max(k * scan, k)]:
            w = self.word(t)
            if (alpha_only and not ALPHA.match(w)) or len(w) < min_len or w.lower() in seen:
                continue
            seen.add(w.lower())
            hits.append(Hit(word=w, score=float(s[t]), token_id=int(t)))
            if len(hits) >= k:
                break
        return hits

    def decode(self, target, k: int = 12, **kw) -> list[AxisDecode]:
        """Both ends of an Axis, of every axis of a Cone, or of a bare direction."""
        from .axes import Axis, Cone
        if isinstance(target, Axis):
            pairs = [(target.name, target.vector)]
        elif isinstance(target, Cone):
            pairs = [(n, target.U[:, i]) for i, n in enumerate(target.names)]
        else:
            pairs = [("direction", np.asarray(target, np.float32))]
        return [AxisDecode(name=n, positive=self.top(v, k=k, sign="+", **kw),
                           negative=self.top(v, k=k, sign="-", **kw)) for n, v in pairs]

    def report(self, target, k: int = 12, **kw) -> str:
        """A printable decode of every axis, both ends."""
        out = [f"{self.spec.key} block {self.block}, {self.n_tokens} tokens"]
        for d in self.decode(target, k=k, **kw):
            out.append(f"  [{d.name:22s}] + {', '.join(d.words('+'))}")
            out.append(f"  {'':24s} - {', '.join(d.words('-'))}")
        return "\n".join(out)

    # ------------------------------------------------------------------ nulls
    def sample_directions(self, n: int = 20000, seed: int = 0) -> np.ndarray:
        """`[n, hidden]` unit transported directions of a random token sample.

        Feeds `purifier`, and gives the null an axis's decode has to be read against: what a
        random direction of this same vocabulary returns.
        """
        rng = np.random.default_rng(seed)
        idx = np.sort(rng.choice(self.n_tokens, size=min(n, self.n_tokens), replace=False))
        pos = {int(t): j for j, t in enumerate(idx)}
        out = np.empty((len(idx), self.spec.hidden_size), np.float32)
        for start, R in self._rows(DECODE_CHUNK):
            want = idx[(idx >= start) & (idx < start + len(R))]
            if not len(want):
                continue
            out[[pos[int(t)] for t in want]] = R[want - start] @ self.J
        return unit(out)

    def purifier(self, k: int = 3, n_rows: int = 20000, seed: int = 0):
        """`(mu, V)` for `purify`, estimated from a token sample of this vocabulary."""
        return common_component(self.sample_directions(n_rows, seed), k=k)


def save_decode(decodes: list[AxisDecode], path: str | Path) -> Path:
    """Write decoded axes to json."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    blob = [{"name": d.name,
             "positive": [{"word": h.word, "score": h.score, "id": h.token_id}
                          for h in d.positive],
             "negative": [{"word": h.word, "score": h.score, "id": h.token_id}
                          for h in d.negative]} for d in decodes]
    path.write_text(json.dumps(blob, indent=1) + "\n")
    return path