# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0

"""``TTWeights``: the device-resident weight set of one (version, embodiment, policy) (IMPLEMENTATION_PLAN §2.2 row
``tt/weights.py``; WP-03).

Pipeline (plan §2.1 row ``build_device_plan``, f1 §6.1, WL §8.2)::

    HostPlanCache(cache_root).load_or_build(version, e, plan_dtype_policy(policy), plan_options(policy, version))
        -> {name: PlanTensor}                                     (torch bf16, host transforms applied, .pt tier)
    per tensor: ttnn.from_torch(t, dtype=dtype_of(t, policy), layout=TILE)     (host tilize / bfp8 pack)
                ttnn.dump_tensor(<dir>/<name>.tensorbin, host_tensor, mode=LOCAL)   (.tensorbin tier, written once)
                ttnn.to_device(host_tensor, device, memory_config=DRAM)             (DRAM interleaved, persistent)
    warm run:   ttnn.load_tensor(<dir>/<name>.tensorbin, device=device)              (no torch, no safetensors)

The ``.tensorbin`` tier lives next to the host tier under the same directory key:
``HostPlanCache.directory(...).parent / "ttnn" / TTPolicy.weights_tag()`` (f1 §6.1 item 1), with a ``manifest.json``
(``common.weights.plan_manifest`` content + the per-tensor device dtype + file facts + the ttnn runtime version) and
a ``.complete`` marker written **last** (v3 semantics of ``models/common/weight_cache.py``): a missing marker, a
manifest identity mismatch, a per-tensor metadata mismatch or a wrong file size => cold path for that tensor.

Measured basis (mb1 §5 Table 5c, 16 x [4096, 8192] = 1.07 GB bf16): cold ``from_torch(device=)`` 0.51 s (host tilize
0.41 s + ``to_device`` 0.068 s at 15.8 GB/s PCIe); ``dump_tensor`` 0.48 s once; warm ``load_tensor(path, device=)``
0.138 s (7.8 GB/s).  bfp8: host pack 0.94 s (2.3x slower than tilize), warm load 0.075 s for 0.57 GB.

Shapes follow the plan (f1 §2): 2-D ``[K, N]`` weights and ``[1, N]`` vectors, all TILE (mb1 §4: ROW_MAJOR gamma /
beta cost 2x and a ``[1,1,1,D]`` ROW_MAJOR gamma is rejected), DRAM interleaved.  ``memory == "HOST"`` entries (N1.7
``vision.pos_embed_table``) are never uploaded; they are kept as torch tensors (``host(name)``) and cached as ``.pt``.

``ttnn`` is imported lazily inside :class:`TTNNBackend` so that CPU tests import this module without a device; the
backend seam also lets CPU tests exercise the cache / manifest / embodiment-switch logic with a fake backend.
"""

from __future__ import annotations

import json
import os
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple

import torch

from models.experimental.gr00t.common import checkpoint as checkpoint_mod
from models.experimental.gr00t.common import weights as W
from models.experimental.gr00t.common.configs import (
    DIT_QUERY_ROWS_PAD,
    HF_SNAPSHOT_SHAS,
    TILE,
    GR00TConfig,
    LayoutConfig,
    get_config,
)
from models.experimental.gr00t.tt.policy import (
    DEVICE_DTYPES,
    TTPolicy,
    dtype_of,
    plan_dtype_policy,
    plan_options,
)

TENSORBIN_FORMAT = "gr00t-ttnn-plan-v1"
TENSORBIN_MARKER = ".complete"
TENSORBIN_MANIFEST = "manifest.json"
TENSORBIN_SUFFIX = ".tensorbin"
HOST_SUFFIX = ".pt"
FORCE_CONVERT_ENV = "GR00T_TT_FORCE_CONVERT"

#: Manifest identity keys (a mismatch in any of them => the whole directory is cold and is rewritten).
_IDENTITY_KEYS: Tuple[str, ...] = (
    "format",
    "version",
    "hf_sha",
    "embodiment_id",
    "plan_dtype_policy",
    "dtype_policy",
    "transform_version",
    "options",
    "arch",
    "runtime",
)
#: Per-tensor metadata that must match for a cached file to be reused.
_TENSOR_KEYS: Tuple[str, ...] = ("category", "logical_shape", "layout", "memory", "plan_dtype", "device_dtype")


# --------------------------------------------------------------------------------------------------------------------
# 1. Upload plan (headers only; CPU-testable)
# --------------------------------------------------------------------------------------------------------------------


@dataclass(frozen=True)
class UploadEntry:
    """One plan tensor with its device dtype decided (``PlanEntry`` + :func:`policy.dtype_of`)."""

    name: str
    category: str
    logical_shape: Tuple[int, ...]
    plan_dtype: str  # dtype the host plan assigns (bf16 | bfp8_b | fp32)
    device_dtype: str  # dtype uploaded under the TTPolicy
    layout: str  # TILE | ROW_MAJOR
    memory: str  # DRAM | HOST

    @property
    def nbytes(self) -> float:
        """Tile-padded device bytes (WL §5 convention; 0 for host-only tensors)."""
        if self.memory == "HOST":
            return 0.0
        return W.device_bytes(self.logical_shape, self.device_dtype, self.layout)

    @property
    def file_name(self) -> str:
        return self.name + (HOST_SUFFIX if self.memory == "HOST" else TENSORBIN_SUFFIX)

    def meta(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "category": self.category,
            "logical_shape": list(self.logical_shape),
            "plan_dtype": self.plan_dtype,
            "device_dtype": self.device_dtype,
            "layout": self.layout,
            "memory": self.memory,
            "device_bytes": self.nbytes,
        }


def upload_plan(
    version: str,
    embodiment_id: Optional[int],
    policy: TTPolicy,
    options: Optional[W.TransformOptions] = None,
    categories: Optional[Iterable[str]] = None,
    names: Optional[Iterable[str]] = None,
) -> Dict[str, UploadEntry]:
    """Every plan tensor of ``version`` with its device dtype under ``policy`` (no weight is read).

    The enumeration uses the host plan policy :func:`plan_dtype_policy` (so the ``mixed`` candidates are flagged) and
    can be restricted by ``categories`` / ``names`` exactly like ``common.weights.build_device_plan`` (unknown values
    raise).  Order: sorted by ``(category, name)`` -- the deterministic upload order.
    """
    opts = options if options is not None else plan_options(policy, version)
    meta = W.enumerate_plan(version, embodiment_id, plan_dtype_policy(policy), opts)
    wanted = set(meta)
    if categories is not None:
        cats = set(categories)
        unknown = cats - set(W.CATEGORIES)
        if unknown:
            raise KeyError(f"unknown categories {sorted(unknown)}")
        wanted = {n for n in wanted if meta[n].category in cats}
    if names is not None:
        req = set(names)
        missing = req - set(meta)
        if missing:
            raise KeyError(f"unknown plan tensor names {sorted(missing)[:5]}")
        wanted &= req
    out: Dict[str, UploadEntry] = {}
    for name in sorted(wanted, key=lambda n: (meta[n].category, n)):
        e = meta[name]
        dd = dtype_of(e, policy)
        if dd not in DEVICE_DTYPES:
            raise ValueError(f"{name}: dtype decision {dd!r} not in {DEVICE_DTYPES}")
        out[name] = UploadEntry(
            name=name,
            category=e.category,
            logical_shape=tuple(int(d) for d in e.shape),
            plan_dtype=e.dtype,
            device_dtype=dd,
            layout=e.layout,
            memory=e.memory,
        )
    return out


def upload_byte_summary(entries: Mapping[str, UploadEntry]) -> Dict[str, Any]:
    """Per-category / per-dtype device bytes of an upload plan (token table and host tensors excluded, WL §6.1)."""
    cats: Dict[str, Dict[str, float]] = {}
    by_dtype: Dict[str, Dict[str, float]] = {}
    total = 0.0
    n = 0
    for e in entries.values():
        b = e.nbytes
        c = cats.setdefault(e.category, {"count": 0, "bytes": 0.0, "mib": 0.0})
        c["count"] += 1
        c["bytes"] += b
        c["mib"] = c["bytes"] / W.MIB
        if e.memory == "HOST" or e.category == "llm_embed":
            continue
        d = by_dtype.setdefault(e.device_dtype, {"count": 0, "bytes": 0.0, "mib": 0.0})
        d["count"] += 1
        d["bytes"] += b
        d["mib"] = d["bytes"] / W.MIB
        total += b
        n += 1
    return {
        "categories": cats,
        "by_dtype": dict(sorted(by_dtype.items())),
        "total_bytes": total,
        "total_mib": total / W.MIB,
        "n_tensors": n,
    }


# --------------------------------------------------------------------------------------------------------------------
# 2. Backend seam (ttnn imported lazily)
# --------------------------------------------------------------------------------------------------------------------


class TTNNBackend:
    """The real backend: thin wrappers over the ``ttnn`` calls of the module docstring."""

    def __init__(self) -> None:
        import ttnn  # lazy: never at module import (plan §2.1 rules)

        self.ttnn = ttnn
        self._dtypes = {"bf16": ttnn.bfloat16, "bfp8_b": ttnn.bfloat8_b, "fp32": ttnn.float32}
        self._layouts = {"TILE": ttnn.TILE_LAYOUT, "ROW_MAJOR": ttnn.ROW_MAJOR_LAYOUT}

    def runtime_tag(self) -> str:
        """``ttnn`` distribution version (carries the git hash, e.g. ``0.75.0rc10.dev1568+g668c2907575``)."""
        try:
            import importlib.metadata as md

            return "ttnn-" + md.version("ttnn")
        except Exception:  # noqa: BLE001 - metadata absent in some source-tree layouts
            v = getattr(self.ttnn, "__version__", None)
            if not v:
                raise RuntimeError("cannot determine the ttnn version for the tensorbin manifest")
            return "ttnn-" + str(v)

    def from_torch_host(self, t: torch.Tensor, dtype: str, layout: str):
        if dtype not in self._dtypes:
            raise ValueError(f"unknown device dtype {dtype!r}")
        if layout not in self._layouts:
            raise ValueError(f"unknown layout {layout!r}")
        return self.ttnn.from_torch(t, dtype=self._dtypes[dtype], layout=self._layouts[layout])

    def dump(self, path: Path, host_tensor) -> None:
        self.ttnn.dump_tensor(str(path), host_tensor, mode=self.ttnn.DumpTensorMode.LOCAL)

    def load(self, path: Path, device):
        return self.ttnn.load_tensor(str(path), device=device)

    def to_device(self, host_tensor, device):
        return self.ttnn.to_device(host_tensor, device, memory_config=self.ttnn.DRAM_MEMORY_CONFIG)

    def to_torch(self, tensor) -> torch.Tensor:
        return self.ttnn.to_torch(tensor)

    def copy_into(self, src, dst) -> None:
        """``ttnn.copy(input_a, input_b)``: copies ``src`` into ``dst`` in place (same shape / dtype / layout)."""
        self.ttnn.copy(src, dst)

    def deallocate(self, tensor) -> None:
        self.ttnn.deallocate(tensor)

    def synchronize(self, device) -> None:
        if device is not None:
            self.ttnn.synchronize_device(device)

    def describe(self, tensor) -> Dict[str, Any]:
        mc = tensor.memory_config()
        return {
            "shape": list(tensor.shape),
            "padded_shape": list(tensor.padded_shape),
            "dtype": str(tensor.dtype),
            "layout": str(tensor.layout),
            "buffer_type": str(mc.buffer_type),
            "memory_layout": str(mc.memory_layout),
        }


# --------------------------------------------------------------------------------------------------------------------
# 3. TTWeights
# --------------------------------------------------------------------------------------------------------------------


def tensorbin_directory(
    host_cache: W.HostPlanCache, version: str, embodiment_id: int, policy: TTPolicy, options: W.TransformOptions
) -> Path:
    """``HostPlanCache.directory(...).parent / "ttnn" / policy.weights_tag()`` (f1 §6.1 item 1)."""
    return (
        host_cache.directory(version, embodiment_id, plan_dtype_policy(policy), options).parent
        / "ttnn"
        / policy.weights_tag()
    )


def hoist_lengths(layout: LayoutConfig) -> Tuple[int, ...]:
    """Padded key lengths of the cross K/V hoist buffers of ``layout`` (plan §2.2 row ``tt/adapter.py``): the padded
    half of ``cross_text_keys`` / ``cross_image_keys`` / ``cross_all_keys`` that are not None, ascending, unique."""
    out = set()
    for keys in (layout.cross_text_keys, layout.cross_image_keys, layout.cross_all_keys):
        if keys is not None:
            padded = int(keys[1])
            if padded <= 0 or padded % TILE:
                raise ValueError(f"{layout.name}: padded key length {padded} is not a positive multiple of {TILE}")
            out.add(padded)
    if not out:
        raise ValueError(f"{layout.name}: no cross key set defined")
    return tuple(sorted(out))


def rope_transformation_matrix(dhead: int = TILE) -> torch.Tensor:
    """``[1,1,32,32]`` transformation matrix of ``ttnn.experimental.rotary_embedding_llama`` (D8): ``+1`` at
    ``(2i, 2i+1)``, ``-1`` at ``(2i+1, 2i)`` (``models/common/tensor_utils.py::get_rot_transformation_mat``)."""
    if dhead != TILE:
        raise ValueError(f"the rotary_embedding_llama transformation matrix is one tile ({TILE}), got {dhead}")
    m = torch.zeros(1, 1, dhead, dhead, dtype=torch.float32)
    m[..., torch.arange(0, dhead, 2), torch.arange(1, dhead, 2)] = 1.0
    m[..., torch.arange(1, dhead, 2), torch.arange(0, dhead, 2)] = -1.0
    return m.to(torch.bfloat16)


def _as_layout(cfg: GR00TConfig, layout: Any) -> LayoutConfig:
    """Accept a ``LayoutConfig``, a layout name, or a ``tt/layout.py`` ``DeviceLayout`` (WP-02: carries
    ``layout_name`` and ``version``; an object exposing a ``.layout`` ``LayoutConfig`` is accepted too) -- an adapter
    so this WP does not import WP-02's class."""
    if isinstance(layout, LayoutConfig):
        return layout
    if isinstance(layout, str):
        return cfg.layout(layout)
    inner = getattr(layout, "layout", None)
    if isinstance(inner, LayoutConfig):
        return inner
    name = getattr(layout, "layout_name", None)
    if isinstance(name, str):
        version = getattr(layout, "version", cfg.version)
        if version != cfg.version:
            raise ValueError(f"DeviceLayout is for {version}, weights are for {cfg.version}")
        return cfg.layout(name)
    raise TypeError(f"expected a LayoutConfig, a layout name or a DeviceLayout (layout_name), got {type(layout)}")


class TTWeights:
    """Device-resident plan tensors of one (version, embodiment, policy); build with :meth:`load`.

    Public API (plan §2.2): ``get(name)``, ``host(name)``, ``names()``, ``host_names()``, ``by_category()``,
    ``embed_rows(input_ids_padded, image_slots)``, ``constants(layout)``, ``switch_embodiment(e)``, ``release()``;
    facts: ``version``, ``embodiment_id``, ``policy``, ``options``, ``plan_dtype_policy``, ``entries`` (name ->
    :class:`UploadEntry`), ``tensorbin_dir``, ``load_stats``, ``device_bytes``.
    """

    def __init__(
        self,
        version: str,
        embodiment_id: int,
        policy: TTPolicy,
        device: Any,
        options: W.TransformOptions,
        host_cache: W.HostPlanCache,
        backend: Any,
    ) -> None:
        self.version = version
        self.cfg: GR00TConfig = get_config(version)
        if not 0 <= int(embodiment_id) < self.cfg.encoders.max_num_embodiments:
            raise IndexError(f"embodiment id {embodiment_id} out of range")
        self.embodiment_id = int(embodiment_id)
        self.policy = policy
        self.device = device
        self.options = options
        self.plan_dtype_policy = plan_dtype_policy(policy)
        self.host_cache = host_cache
        self.backend = backend
        self.tensorbin_dir: Path = tensorbin_directory(host_cache, version, self.embodiment_id, policy, options)
        self.entries: Dict[str, UploadEntry] = {}
        self._tensors: Dict[str, Any] = {}
        self._host: Dict[str, torch.Tensor] = {}
        self._constants: Dict[str, Dict[str, Any]] = {}
        self._embed_table: Optional[torch.Tensor] = None
        self.load_stats: Dict[str, Any] = {}

    # ------------------------------------------------------------------ construction
    @classmethod
    def load(
        cls,
        version: str,
        embodiment_id: Optional[int],
        policy: TTPolicy,
        device: Any,
        cache_root: Optional[Path] = None,
        categories: Optional[Iterable[str]] = None,
        names: Optional[Iterable[str]] = None,
        *,
        layout_name: Optional[str] = None,
        backend: Any = None,
        force_convert: bool = False,
    ) -> "TTWeights":
        """Upload (or warm-load) the plan tensors of ``version`` / ``embodiment_id`` under ``policy`` to ``device``.

        * ``embodiment_id=None`` -> the config's canonical embodiment.
        * ``categories`` / ``names`` restrict the upload (tests, partial loads); the ``.tensorbin`` tier is per tensor
          so a later full load reuses whatever a partial load wrote.
        * ``layout_name`` selects the layout whose ``S_pad`` decides ``gate_up_layout="auto"`` (default canonical).
        * ``force_convert`` (or ``GR00T_TT_FORCE_CONVERT=1``) ignores both cache tiers.
        * ``backend`` defaults to :class:`TTNNBackend`; CPU tests inject a fake.
        """
        cfg = get_config(version)
        e = cfg.canonical_embodiment_id if embodiment_id is None else int(embodiment_id)
        if not isinstance(policy, TTPolicy):
            raise TypeError(f"policy must be a TTPolicy, got {type(policy).__name__}")
        opts = plan_options(policy, version, layout_name)
        host_cache = W.HostPlanCache(cache_root)
        be = backend if backend is not None else TTNNBackend()
        self = cls(version, e, policy, device, opts, host_cache, be)
        force = bool(force_convert) or os.environ.get(FORCE_CONVERT_ENV, "0") == "1"
        self._load(categories, names, force)
        return self

    def _load(self, categories: Optional[Iterable[str]], names: Optional[Iterable[str]], force: bool) -> None:
        t_start = time.perf_counter()
        full_request = categories is None and names is None
        full_entries = upload_plan(self.version, self.embodiment_id, self.policy, self.options)
        wanted = (
            full_entries
            if full_request
            else upload_plan(self.version, self.embodiment_id, self.policy, self.options, categories, names)
        )
        if not wanted:
            raise ValueError("the requested categories / names select no plan tensor")
        self.entries = dict(wanted)
        expected_identity = self._identity()
        manifest = None if force else self._read_manifest(expected_identity)
        warm: List[str] = []
        cold: List[str] = []
        for name, entry in wanted.items():
            if manifest is not None and self._file_is_warm(manifest, entry):
                warm.append(name)
            else:
                cold.append(name)

        # --- warm tensors: .tensorbin straight to DRAM (no torch, no safetensors)
        t0 = time.perf_counter()
        for name in warm:
            entry = wanted[name]
            path = self.tensorbin_dir / entry.file_name
            if entry.memory == "HOST":
                self._host[name] = torch.load(path, map_location="cpu", weights_only=True)
            else:
                self._tensors[name] = self.backend.load(path, self.device)
        self.backend.synchronize(self.device)
        t_warm = time.perf_counter() - t0

        # --- cold tensors: host plan -> tilize/pack -> dump -> upload
        t_host = 0.0
        t_cold = 0.0
        file_facts: Dict[str, Dict[str, Any]] = {}
        if cold:
            t0 = time.perf_counter()
            plan = self._host_plan(self.embodiment_id, cold, full_request, force)
            t_host = time.perf_counter() - t0
            t0 = time.perf_counter()
            self.tensorbin_dir.mkdir(parents=True, exist_ok=True)
            for name in cold:
                entry = wanted[name]
                pt = plan[name]
                self._check_plan_tensor(entry, pt)
                path = self.tensorbin_dir / entry.file_name
                if entry.memory == "HOST":
                    torch.save(pt.tensor.contiguous(), path)
                    self._host[name] = pt.tensor
                else:
                    host_t = self.backend.from_torch_host(pt.tensor, entry.device_dtype, entry.layout)
                    self.backend.dump(path, host_t)
                    self._tensors[name] = self.backend.to_device(host_t, self.device)
                file_facts[name] = {"file": path.name, "file_bytes": int(path.stat().st_size)}
            self.backend.synchronize(self.device)
            t_cold = time.perf_counter() - t0
            self._write_manifest(full_entries, expected_identity, manifest, file_facts)

        n_dev = sum(1 for n in wanted if wanted[n].memory != "HOST")
        self.load_stats = {
            "path": "warm" if not cold else ("cold" if not warm else "partial"),
            "seconds": time.perf_counter() - t_start,
            "warm_load_s": t_warm,
            "host_plan_s": t_host,
            "convert_dump_upload_s": t_cold,
            "n_warm": len(warm),
            "n_cold": len(cold),
            "n_device_tensors": n_dev,
            "n_host_tensors": len(wanted) - n_dev,
            "device_bytes": self.device_bytes,
            "tensorbin_dir": str(self.tensorbin_dir),
            "runtime": expected_identity["runtime"],
            "forced": bool(force),
        }

    def _host_plan(self, embodiment_id: int, names: Sequence[str], full_request: bool, force: bool) -> W.DevicePlan:
        """Host plan tensors for ``names`` of embodiment ``embodiment_id``: the complete ``.pt`` tier when present,
        else ``load_or_build`` for a full request (which populates the ``.pt`` tier), else a partial build."""
        d = self.host_cache.directory(self.version, embodiment_id, self.plan_dtype_policy, self.options)
        if not force and self.host_cache.is_complete(d):
            return self.host_cache.load(d, names=names)
        if full_request:
            return self.host_cache.load_or_build(
                self.version, embodiment_id, self.plan_dtype_policy, self.options, force_convert=force
            )
        return W.build_device_plan(
            self.version, embodiment_id, self.plan_dtype_policy, options=self.options, names=names
        )

    @staticmethod
    def _check_plan_tensor(entry: UploadEntry, pt: W.PlanTensor) -> None:
        if tuple(pt.tensor.shape) != entry.logical_shape:
            raise ValueError(f"{entry.name}: host tensor shape {tuple(pt.tensor.shape)} != plan {entry.logical_shape}")
        if pt.dtype != entry.plan_dtype or pt.layout != entry.layout or pt.memory != entry.memory:
            raise ValueError(
                f"{entry.name}: host plan metadata ({pt.dtype}, {pt.layout}, {pt.memory}) != enumerated "
                f"({entry.plan_dtype}, {entry.layout}, {entry.memory})"
            )
        want_torch = torch.float32 if entry.plan_dtype == "fp32" else torch.bfloat16
        if pt.tensor.dtype != want_torch:
            raise TypeError(f"{entry.name}: host tensor is {pt.tensor.dtype}, expected {want_torch}")

    # ------------------------------------------------------------------ manifest / marker
    def _identity(self) -> Dict[str, Any]:
        from dataclasses import asdict

        return {
            "format": TENSORBIN_FORMAT,
            "version": self.version,
            "hf_sha": HF_SNAPSHOT_SHAS[self.version],
            "embodiment_id": self.embodiment_id,
            "plan_dtype_policy": self.plan_dtype_policy,
            "dtype_policy": self.policy.dtype_policy,
            "transform_version": W.TRANSFORM_VERSION,
            "options": asdict(self.options),
            "arch": self.host_cache.arch,
            "runtime": self.backend.runtime_tag(),
        }

    def is_complete(self) -> bool:
        marker = self.tensorbin_dir / TENSORBIN_MARKER
        if not marker.is_file() or not (self.tensorbin_dir / TENSORBIN_MANIFEST).is_file():
            return False
        try:
            return marker.read_text().strip() == TENSORBIN_FORMAT
        except OSError:
            return False

    def _read_manifest(self, expected_identity: Mapping[str, Any]) -> Optional[Dict[str, Any]]:
        """The manifest when the marker is present and the identity matches; otherwise None (cold directory)."""
        if not self.is_complete():
            return None
        try:
            with open(self.tensorbin_dir / TENSORBIN_MANIFEST) as fh:
                m = json.load(fh)
        except (OSError, ValueError):
            return None
        if any(m.get(k) != expected_identity[k] for k in _IDENTITY_KEYS):
            return None
        if not isinstance(m.get("tensors"), dict):
            return None
        return m

    def _file_is_warm(self, manifest: Mapping[str, Any], entry: UploadEntry) -> bool:
        rec = manifest["tensors"].get(entry.name)
        if rec is None or rec.get("file_bytes") is None:
            return False
        meta = entry.meta()
        if any(rec.get(k) != meta[k] for k in _TENSOR_KEYS):
            return False
        path = self.tensorbin_dir / entry.file_name
        if rec.get("file") != path.name:
            return False
        return path.is_file() and path.stat().st_size == int(rec["file_bytes"])

    def _write_manifest(
        self,
        full_entries: Mapping[str, UploadEntry],
        identity: Mapping[str, Any],
        previous: Optional[Mapping[str, Any]],
        file_facts: Mapping[str, Mapping[str, Any]],
    ) -> Dict[str, Any]:
        """Rewrite the manifest for the whole plan (file facts for the tensors on disk) then the marker, last.

        With no valid previous manifest the directory is reset (stale ``.tensorbin`` / ``.pt`` files of a different
        identity are removed) so a marker never outlives files it does not describe (WL §8.2 item 2)."""
        directory = self.tensorbin_dir
        directory.mkdir(parents=True, exist_ok=True)
        marker = directory / TENSORBIN_MARKER
        if marker.exists():
            marker.unlink()
        if previous is None:
            keep = {full_entries[n].file_name for n in file_facts}
            for p in directory.iterdir():
                if p.is_file() and p.suffix in (TENSORBIN_SUFFIX, HOST_SUFFIX) and p.name not in keep:
                    p.unlink()
        tensors: Dict[str, Any] = {}
        total = 0
        n_files = 0
        for name, entry in full_entries.items():
            rec = entry.meta()
            rec["file"] = entry.file_name
            rec["file_bytes"] = None
            if name in file_facts:
                rec["file_bytes"] = int(file_facts[name]["file_bytes"])
            elif previous is not None and self._file_is_warm(previous, entry):
                rec["file_bytes"] = int(previous["tensors"][name]["file_bytes"])
            if rec["file_bytes"] is not None:
                total += rec["file_bytes"]
                n_files += 1
            tensors[name] = rec
        manifest: Dict[str, Any] = dict(identity)
        manifest.update(
            {
                "policy": self.policy.to_dict(),
                "n_tensors": len(tensors),
                "n_files": n_files,
                "file_bytes_total": total,
                "device_bytes_total": upload_byte_summary(full_entries)["total_bytes"],
                "complete": n_files == len(tensors),
                "tensors": tensors,
            }
        )
        with open(directory / TENSORBIN_MANIFEST, "w") as fh:
            json.dump(manifest, fh, indent=1, sort_keys=True)
        marker.write_text(TENSORBIN_FORMAT + "\n")
        return manifest

    def manifest(self) -> Dict[str, Any]:
        """The manifest on disk (raises when the directory is not complete)."""
        if not self.is_complete():
            raise FileNotFoundError(f"{self.tensorbin_dir}: no complete tensorbin tier")
        with open(self.tensorbin_dir / TENSORBIN_MANIFEST) as fh:
            return json.load(fh)

    # ------------------------------------------------------------------ access
    def get(self, name: str):
        """Device tensor ``name`` (``KeyError`` for unknown names; host-only entries point to :meth:`host`)."""
        try:
            return self._tensors[name]
        except KeyError:
            if name in self._host:
                raise KeyError(f"{name!r} is a host-only plan tensor; use host({name!r})") from None
            raise KeyError(f"unknown or not-loaded plan tensor {name!r}") from None

    def host(self, name: str) -> torch.Tensor:
        try:
            return self._host[name]
        except KeyError:
            raise KeyError(f"unknown or not-loaded host plan tensor {name!r}") from None

    def __contains__(self, name: object) -> bool:
        return name in self._tensors

    def __len__(self) -> int:
        return len(self._tensors)

    def names(self) -> List[str]:
        return sorted(self._tensors)

    def host_names(self) -> List[str]:
        return sorted(self._host)

    def by_category(self) -> Dict[str, List[str]]:
        out: Dict[str, List[str]] = {}
        for name in self.names():
            out.setdefault(self.entries[name].category, []).append(name)
        return out

    @property
    def device_bytes(self) -> float:
        return float(sum(self.entries[n].nbytes for n in self._tensors))

    def dtype_histogram(self) -> Dict[str, int]:
        out: Dict[str, int] = {}
        for n in self._tensors:
            d = self.entries[n].device_dtype
            out[d] = out.get(d, 0) + 1
        return dict(sorted(out.items()))

    # ------------------------------------------------------------------ host gather of token rows
    def _embed_table_tensor(self) -> torch.Tensor:
        """The bf16 ``embed_tokens`` table ``[V, 2048]``, read lazily from the checkpoint on first use (WL §4.4 option
        A: host gather; the resident-table + ``ttnn.embedding`` alternative is the WP-P lever)."""
        if self._embed_table is None:
            ck = checkpoint_mod.LazyCheckpoint(self.version)
            key = self.cfg.checkpoint.llm_prefix + "embed_tokens.weight"
            table = ck.get(key)
            want = (self.cfg.llm.vocab, self.cfg.llm.hidden)
            if tuple(table.shape) != want or table.dtype != torch.bfloat16:
                raise ValueError(f"{key}: shape/dtype {tuple(table.shape)}/{table.dtype} != {want}/bf16")
            self._embed_table = table
        return self._embed_table

    def embed_rows(self, input_ids_padded: torch.Tensor, image_slots: Any = None) -> torch.Tensor:
        """Token-embedding rows of a padded prompt: ``[1,1,S_pad,2048]`` bf16 (plan §3.1 host per call).

        ``input_ids_padded`` is ``SlotLayout.input_ids_padded`` ``(1, S_pad)``; ``image_slots`` is a ``SlotLayout`` /
        ``LayoutConfig`` (their ``.image_slots``) or a sequence of half-open ``(lo, hi)`` runs whose rows are zeroed
        (they are replaced by vision rows on device, ``TokenSplicer``); ``None`` zeroes nothing.  Exact gather
        (f1 §0: embed-row gather exact 3/3).
        """
        ids = torch.as_tensor(input_ids_padded)
        if ids.ndim != 2 or ids.shape[0] != 1:
            raise ValueError(f"input_ids_padded must be (1, S_pad), got {tuple(ids.shape)}")
        s_pad = int(ids.shape[1])
        if s_pad % TILE:
            raise ValueError(f"S_pad={s_pad} is not a multiple of {TILE}")
        table = self._embed_table_tensor()
        rows = W.embed_rows(table, ids.long()).clone()  # [1, S_pad, D], owns its bytes
        slots: Tuple[Tuple[int, int], ...]
        if image_slots is None:
            slots = ()
        elif hasattr(image_slots, "image_slots"):
            slots = tuple(tuple(int(x) for x in run) for run in image_slots.image_slots)
        else:
            slots = tuple(tuple(int(x) for x in run) for run in image_slots)
        for lo, hi in slots:
            if not 0 <= lo < hi <= s_pad:
                raise ValueError(f"image slot ({lo}, {hi}) outside [0, {s_pad})")
            rows[:, lo:hi] = 0
        return rows.reshape(1, 1, s_pad, table.shape[1]).to(torch.bfloat16)

    # ------------------------------------------------------------------ constants
    def _action_rows(self, layout: LayoutConfig) -> Tuple[int, int]:
        """``[M-H, M)``: the rows of the logical-64 DiT stream that hold the action tokens (plan §2.2 rows
        ``tt/encoders.py``: ``noise_to_xt`` / ``b3_pos_rows`` / Euler row mask)."""
        m = self.cfg.dit.n_query_tokens
        h = self.cfg.encoders.action_horizon
        if (
            layout.dit_m_logical != m
            or layout.dit_m_pad != DIT_QUERY_ROWS_PAD
            or self.cfg.dit.m_pad != DIT_QUERY_ROWS_PAD
        ):
            raise ValueError(
                f"{layout.name}: dit_m_logical/dit_m_pad {layout.dit_m_logical}/{layout.dit_m_pad} != {m}/{DIT_QUERY_ROWS_PAD}"
            )
        if not 0 < h <= m:
            raise ValueError(f"action horizon {h} does not fit in {m} query tokens")
        return m - h, m

    def b3_pos_rows_host(self, layout: LayoutConfig, b3_pos: Optional[torch.Tensor] = None) -> torch.Tensor:
        """``[1,1,64,1536]``: ``enc.action.b3_pos`` (``b3 + position_embedding[0:H]``, f1 §2) at rows ``M-H..M``,
        zeros elsewhere, in the precompute dtype."""
        lo, hi = self._action_rows(layout)
        src = b3_pos if b3_pos is not None else self.backend.to_torch(self.get("enc.action.b3_pos"))
        emb = self.cfg.encoders.input_embedding_dim
        if tuple(src.shape) != (hi - lo, emb):
            raise ValueError(f"enc.action.b3_pos shape {tuple(src.shape)} != {(hi - lo, emb)}")
        want = torch.float32 if self.policy.precompute_dtype == "fp32" else torch.bfloat16
        out = torch.zeros(1, 1, DIT_QUERY_ROWS_PAD, emb, dtype=want)
        out[0, 0, lo:hi] = src.to(want)
        return out

    def euler_row_mask_host(self, layout: LayoutConfig) -> torch.Tensor:
        """``[1,1,64,A_pad]`` bf16: 1 on the action rows ``M-H..M`` (``cfg.dit.velocity_rows`` of the decoder output),
        0 on the state / future-token / pad rows, so ``x_t += dt * (v * mask)`` runs on all 64 rows (plan §2.2 row
        ``tt/encoders.py`` ``EulerUpdate``)."""
        lo, hi = self._action_rows(layout)
        vr = self.cfg.dit.velocity_rows
        if vr.start != lo or vr.stop is not None:
            raise ValueError(f"cfg.dit.velocity_rows {vr} disagrees with the action rows [{lo}, {hi})")
        out = torch.zeros(1, 1, DIT_QUERY_ROWS_PAD, self.cfg.encoders.action_dim_pad, dtype=torch.bfloat16)
        out[0, 0, lo:hi] = 1.0
        return out

    def constants(self, layout: Any) -> Dict[str, Any]:
        """Constant device tensors of one static layout (bf16 TILE DRAM), built once per layout name and cached:

        * ``trans_mat`` ``[1,1,32,32]`` for ``rotary_embedding_llama`` (:func:`rope_transformation_matrix`, D8);
        * ``q_dummy_<L_pad>`` ``[1,1,L_pad,2048]`` zeros per hoist length (D7: ``nlp_create_qkv_heads(q_dummy,
          input_kv=kv)`` needs an equal-S Q; 2048 = 32 heads x 64 padded head dim);
        * ``euler_row_mask`` ``[1,1,64,A_pad]`` (:meth:`euler_row_mask_host`);
        * ``b3_pos_rows`` ``[1,1,64,1536]`` (:meth:`b3_pos_rows_host`; refreshed in place by :meth:`switch_embodiment`).

        ``layout`` is a ``LayoutConfig``, its name, or an object with a ``.layout`` ``LayoutConfig`` (``DeviceLayout``).
        The attention masks belong to ``tt/layout.py`` (WP-02) and are not built here.
        """
        lay = _as_layout(self.cfg, layout)
        if lay.name in self._constants:
            return self._constants[lay.name]
        inner = self.cfg.dit.n_heads * self.cfg.dit.head_dim_pad
        host: Dict[str, torch.Tensor] = {"trans_mat": rope_transformation_matrix()}
        for l_pad in hoist_lengths(lay):
            host[f"q_dummy_{l_pad}"] = torch.zeros(1, 1, l_pad, inner, dtype=torch.bfloat16)
        host["euler_row_mask"] = self.euler_row_mask_host(lay)
        host["b3_pos_rows"] = self.b3_pos_rows_host(lay)
        out: Dict[str, Any] = {}
        for key, t in host.items():
            dtype = "fp32" if t.dtype == torch.float32 else "bf16"
            out[key] = self.backend.to_device(self.backend.from_torch_host(t, dtype, "TILE"), self.device)
        self.backend.synchronize(self.device)
        self._constants[lay.name] = out
        return out

    # ------------------------------------------------------------------ embodiment switch
    def encoder_names(self) -> List[str]:
        return [n for n in self.names() if self.entries[n].category == "encoders"]

    def switch_embodiment(self, embodiment_id: int) -> Dict[str, Any]:
        """Re-point the resident encoder buffers at embodiment ``embodiment_id`` **in place** (``ttnn.copy`` into the
        existing DRAM buffers, so trace addresses stay valid; WL §8.2 item 5).

        Only category ``encoders`` is rebuilt (the ``.pt`` tier of the new slot when complete, else a partial build
        of that category); every shape / dtype must match the resident buffer (raise otherwise).  ``b3_pos_rows`` of
        every materialised :meth:`constants` set is refreshed the same way.  Returns the switch statistics.
        """
        e = int(embodiment_id)
        if not 0 <= e < self.cfg.encoders.max_num_embodiments:
            raise IndexError(f"embodiment id {e} out of range [0, {self.cfg.encoders.max_num_embodiments})")
        enc = self.encoder_names()
        full_enc = set(
            upload_plan(self.version, self.embodiment_id, self.policy, self.options, categories=("encoders",))
        )
        if not enc or set(enc) != full_enc:
            raise RuntimeError(
                f"switch_embodiment needs the complete 'encoders' category resident ({len(enc)}/{len(full_enc)} "
                "tensors loaded); load with categories including 'encoders'"
            )
        t0 = time.perf_counter()
        plan = self._host_plan(e, enc, full_request=False, force=False)
        missing = set(enc) - set(plan)
        if missing:
            raise KeyError(f"encoder plan of embodiment {e} lacks {sorted(missing)[:5]}")
        t_host = time.perf_counter() - t0
        t0 = time.perf_counter()
        n_bytes = 0.0
        for name in enc:
            entry = self.entries[name]
            pt = plan[name]
            self._check_plan_tensor(entry, pt)
            staged = self.backend.to_device(
                self.backend.from_torch_host(pt.tensor, entry.device_dtype, entry.layout), self.device
            )
            self.backend.copy_into(staged, self._tensors[name])
            self.backend.deallocate(staged)
            n_bytes += entry.nbytes
        b3 = plan["enc.action.b3_pos"].tensor if "enc.action.b3_pos" in plan else None
        for lay_name, consts in self._constants.items():
            if "b3_pos_rows" in consts:
                if b3 is None:
                    raise RuntimeError("enc.action.b3_pos not resident; cannot refresh b3_pos_rows")
                host_t = self.b3_pos_rows_host(self.cfg.layout(lay_name), b3)
                dtype = "fp32" if host_t.dtype == torch.float32 else "bf16"
                staged = self.backend.to_device(self.backend.from_torch_host(host_t, dtype, "TILE"), self.device)
                self.backend.copy_into(staged, consts["b3_pos_rows"])
                self.backend.deallocate(staged)
        self.backend.synchronize(self.device)
        previous = self.embodiment_id
        self.embodiment_id = e
        return {
            "from": previous,
            "to": e,
            "n_tensors": len(enc),
            "device_bytes": n_bytes,
            "host_plan_s": t_host,
            "copy_s": time.perf_counter() - t0,
        }

    # ------------------------------------------------------------------ teardown
    def release(self) -> None:
        """Deallocate every device tensor (plan tensors and constants)."""
        for consts in self._constants.values():
            for t in consts.values():
                self.backend.deallocate(t)
        self._constants.clear()
        for t in self._tensors.values():
            self.backend.deallocate(t)
        self._tensors.clear()
        self._host.clear()
        self._embed_table = None


__all__ = [
    "TENSORBIN_FORMAT",
    "TENSORBIN_MARKER",
    "TENSORBIN_MANIFEST",
    "TENSORBIN_SUFFIX",
    "UploadEntry",
    "upload_plan",
    "upload_byte_summary",
    "TTNNBackend",
    "tensorbin_directory",
    "hoist_lengths",
    "rope_transformation_matrix",
    "TTWeights",
]