"""Registry of models a U-Space basis can be built for, with their pinned revisions.

The readout block is the model's block at relative depth 2/3. It is recorded explicitly so a
basis is reproducible without fetching a config, and `tests/test_models.py` checks every
entry against `rel_to_block`.
"""
from __future__ import annotations

from dataclasses import dataclass, field

#: relative depth the state is READ at, one of three pre-registered depths (2/3, 0.75, 0.80)
READOUT_REL = 2 / 3

#: relative depth the axes DECODE at. Not the same depth: the axes return recognisable words
#: only around rel 0.75-0.80, while the score is best at 2/3. Reading an axis's identity and
#: reading its activation are therefore two different depths, and both are reported as such.
DECODE_REL = 0.75

#: the published Jacobian-lens collection
LENS_REPO = "neuronpedia/jacobian-lens"
LENS_REV = "0731326edff4ae730ffc5356fe1a4728c748b3a6"


@dataclass(frozen=True)
class LensFile:
    """Where a model's Jacobian lens file lives.

    `repo`/`file` name a file on the Hub; `path` a local `.pt` (all fitted blocks) or `.npy`
    (one block, already extracted).
    """
    repo: str | None = None
    file: str | None = None
    revision: str | None = None
    path: str | None = None

    def __post_init__(self) -> None:
        if not self.path and not (self.repo and self.file):
            raise ValueError("a LensFile needs either `path` or both `repo` and `file`")


@dataclass(frozen=True)
class ModelSpec:
    """A model, its pinned revision, its lens, and the block the state is read at."""
    key: str
    hf: str
    revision: str
    n_blocks: int
    hidden_size: int
    readout_block: int
    lens: LensFile
    think_end: str
    think_start: str = "<think>"
    #: the sampling settings the reported runs used
    sampling: dict = field(default_factory=lambda: {"temperature": 1.0, "top_p": 0.95,
                                                    "top_k": 20})
    aliases: tuple[str, ...] = field(default=())
    paper_model: bool = False


def rel_to_block(rel: float, n_blocks: int) -> int:
    """Relative depth -> lens block key.

    The lens is keyed by decoder block `b` and transports that block's *output*, i.e. hidden
    state index `b + 1`, so `rel = (b + 1) / n_blocks`.
    """
    h = round(rel * n_blocks)
    return min(max(h, 1), n_blocks - 1) - 1


MODELS: dict[str, ModelSpec] = {}


def _add(spec: ModelSpec) -> None:
    MODELS[spec.key] = spec


_add(ModelSpec(
    key="gemma-4-31b-it",
    hf="google/gemma-4-31B-it",
    revision="842da3794eaa0b77d5f08bae87a17459d91ff475",
    n_blocks=60, hidden_size=5376, readout_block=39,
    lens=LensFile(repo=LENS_REPO, revision=LENS_REV,
              file="gemma-4-31b/jlens/Salesforce-wikitext/gemma-4-31B_jacobian_lens.pt"),
    think_end="<channel|>", think_start="<think>",
    sampling={"temperature": 1.0, "top_p": 0.95, "top_k": 64},
    aliases=("gemma",), paper_model=True,
))

_add(ModelSpec(
    key="qwen3.5-27b",
    hf="Qwen/Qwen3.5-27B",
    revision="fc05daec18b0a78c049392ed2e771dde82bdf654",
    n_blocks=64, hidden_size=5120, readout_block=42,
    lens=LensFile(repo=LENS_REPO, revision=LENS_REV,
              file="qwen3.5-27b/jlens/Salesforce-wikitext/Qwen3.5-27B_jacobian_lens.pt"),
    think_end="</think>", aliases=("qwen",), paper_model=True,
))

# No lens is published for Mistral, so this one was fitted for the paper and mirrored.
_add(ModelSpec(
    key="magistral-small-2507",
    hf="mistralai/Magistral-Small-2507",
    revision="b3583a426ca23186681c4297926dfe9219be6161",
    n_blocks=40, hidden_size=5120, readout_block=26,
    lens=LensFile(repo="loose-bits/uspace-lenses",
              revision="59e7f1c39e6a6379321b55ac395f5d7ed23cebcb",
              file="magistral_100fix_jacobian_lens.pt"),
    think_end="[/THINK]", think_start="[THINK]",
    sampling={"temperature": 0.7, "top_p": 0.95},
    aliases=("magistral",), paper_model=True,
))

# Not a paper model: small enough to run interactively, and its published lens is fitted on
# more prompts (1000) than any of the three above.
_add(ModelSpec(
    key="qwen3.5-4b",
    hf="Qwen/Qwen3.5-4B",
    revision="851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a",
    n_blocks=32, hidden_size=2560, readout_block=20,
    lens=LensFile(repo=LENS_REPO, revision=LENS_REV,
              file="qwen3.5-4b/jlens/Salesforce-wikitext/Qwen3.5-4B_jacobian_lens_n1000.pt"),
    think_end="</think>", aliases=("qwen4b",),
))

_ALIASES = {a: s.key for s in MODELS.values() for a in s.aliases}


def get(model: str | ModelSpec) -> ModelSpec:
    """Look a model up by key or alias."""
    if isinstance(model, ModelSpec):
        return model
    key = _ALIASES.get(model, model)
    if key not in MODELS:
        known = ", ".join(sorted(MODELS) + sorted(_ALIASES))
        raise KeyError(f"unknown model {model!r}; known: {known}")
    return MODELS[key]


def keys(paper_only: bool = False) -> list[str]:
    """Registry keys, optionally only the three models the paper reports."""
    return sorted(k for k, s in MODELS.items() if s.paper_model or not paper_only)


def decode_block(model: str | ModelSpec) -> int:
    """The block to read an axis's *identity* at, which is not its readout block."""
    spec = get(model)
    return rel_to_block(DECODE_REL, spec.n_blocks)