# SPDX-License-Identifier: Apache-2.0
"""C04: weights -- ONNX initializers addressed by the node that consumes them, BN folding in float64, safetensors /
``.pth`` loading, and an on-disk cache of prepared (tilized, converted) device tensors.

Why by consuming node: initializer names drift between exports of the same network (FRNet), and MatMul weights are
often anonymous (``onnx::MatMul_1234`` in Diffusion Planner), while node names (``/backbone/conv1/Conv``) are stable.
``OnnxWeights`` parses the graph as data only (no onnxruntime, no execution); external-data paths must stay inside
the model's directory.

BN folding: every fold is computed in float64 and rounded once at the end (default float32; pass ``dtype=None`` to
keep float64 and round straight to the device dtype later).

The device cache stores ``ttnn`` host tensors (``.tensorbin``) under ``$TT_CACHE_PATH`` (the tt-model container
mounts it at ``/tensor-cache``) or ``~/.cache/ttaw``, keyed by a digest of the source weights, so a new weights
revision never reuses stale tensors. ``TTAW_WEIGHT_CACHE=0`` disables it (A/B switch, read when the cache is built).

Only ``numpy`` is imported at module level; ``onnx``, ``safetensors``, ``torch`` and ``ttnn`` are imported by the
functions that need them.
"""
from __future__ import annotations

import fnmatch
import hashlib
import os
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable, Dict, Iterator, List, Mapping, Optional, Tuple, Union

import numpy as np

__all__ = [
    "BatchNorm",
    "fold_bn",
    "fold_bn_conv",
    "fold_bn_linear",
    "fold_bn_conv_transpose",
    "OnnxNode",
    "ConvParams",
    "GemmParams",
    "OnnxWeights",
    "StateDict",
    "load_safetensors",
    "load_torch_checkpoint",
    "load_state_dict",
    "file_sha256",
    "cache_root",
    "WeightCache",
]

PathLike = Union[str, os.PathLike]


# ---------------------------------------------------------------------------------------------- BN folding

@dataclass(frozen=True)
class BatchNorm:
    """Inference BatchNorm parameters (any of BN1d / BN2d / ONNX BatchNormalization), stored in float64."""

    gamma: np.ndarray
    beta: np.ndarray
    mean: np.ndarray
    var: np.ndarray
    eps: float = 1e-5

    def __post_init__(self) -> None:
        for name in ("gamma", "beta", "mean", "var"):
            object.__setattr__(self, name, np.asarray(getattr(self, name), dtype=np.float64).reshape(-1))
        n = self.gamma.shape[0]
        if not (self.beta.shape[0] == self.mean.shape[0] == self.var.shape[0] == n):
            raise ValueError("BatchNorm parameter lengths differ")

    @property
    def channels(self) -> int:
        return int(self.gamma.shape[0])

    @property
    def scale(self) -> np.ndarray:
        """``gamma / sqrt(var + eps)`` (float64)."""
        return self.gamma / np.sqrt(self.var + self.eps)

    @property
    def shift(self) -> np.ndarray:
        """``beta - mean * scale`` (float64): ``bn(x) == x * scale + shift``."""
        return self.beta - self.mean * self.scale

    @classmethod
    def from_state_dict(cls, sd: Mapping[str, Any], prefix: str, eps: float = 1e-5) -> "BatchNorm":
        """torch naming: ``<prefix>.weight``, ``.bias``, ``.running_mean``, ``.running_var``."""
        p = prefix.rstrip(".") + "." if prefix else ""
        return cls(np.asarray(sd[p + "weight"]), np.asarray(sd[p + "bias"]), np.asarray(sd[p + "running_mean"]),
                   np.asarray(sd[p + "running_var"]), eps)

    def apply(self, x: np.ndarray, axis: int = 1) -> np.ndarray:
        """Reference application in float64 (``axis`` = channel axis), for tests."""
        shape = [1] * np.ndim(x)
        shape[axis] = -1
        return np.asarray(x, np.float64) * self.scale.reshape(shape) + self.shift.reshape(shape)


def _finish(w: np.ndarray, b: np.ndarray, dtype: Optional[Any]) -> Tuple[np.ndarray, np.ndarray]:
    if dtype is None:
        return w, b
    return w.astype(dtype), b.astype(dtype)


def fold_bn(weight: np.ndarray, bias: Optional[np.ndarray], bn: BatchNorm, *, axis: int = 0,
            dtype: Optional[Any] = np.float32) -> Tuple[np.ndarray, np.ndarray]:
    """Fold ``bn`` (applied to the layer's output) into ``weight`` / ``bias``: ``W' = W * s`` along the output axis
    ``axis`` and ``b' = (b - mean) * s + beta``, all in float64, one final rounding to ``dtype``."""
    w = np.asarray(weight, np.float64)
    if w.shape[axis] != bn.channels:
        raise ValueError(f"weight axis {axis} has {w.shape[axis]} channels, BatchNorm has {bn.channels}")
    shape = [1] * w.ndim
    shape[axis] = -1
    s = bn.scale
    b = np.zeros(bn.channels, np.float64) if bias is None else np.asarray(bias, np.float64).reshape(-1)
    return _finish(w * s.reshape(shape), (b - bn.mean) * s + bn.beta, dtype)


def fold_bn_conv(weight: np.ndarray, bias: Optional[np.ndarray], bn: BatchNorm, *,
                 dtype: Optional[Any] = np.float32) -> Tuple[np.ndarray, np.ndarray]:
    """Conv1d/2d/3d weight ``[Cout, Cin/groups, k...]`` followed by BN (any ``groups``)."""
    return fold_bn(weight, bias, bn, axis=0, dtype=dtype)


def fold_bn_linear(weight: np.ndarray, bias: Optional[np.ndarray], bn: BatchNorm, *, layout: str = "out_in",
                   dtype: Optional[Any] = np.float32) -> Tuple[np.ndarray, np.ndarray]:
    """Linear / Gemm / MatMul followed by BN1d. ``layout="out_in"``: torch ``Linear`` and ONNX Gemm with
    ``transB=1`` (``[out, in]``); ``"in_out"``: ONNX MatMul / Gemm ``transB=0`` (``[in, out]``)."""
    if layout not in ("out_in", "in_out"):
        raise ValueError("layout must be 'out_in' or 'in_out'")
    return fold_bn(weight, bias, bn, axis=0 if layout == "out_in" else -1, dtype=dtype)


def fold_bn_conv_transpose(weight: np.ndarray, bias: Optional[np.ndarray], bn: BatchNorm, *, groups: int = 1,
                           dtype: Optional[Any] = np.float32) -> Tuple[np.ndarray, np.ndarray]:
    """ConvTranspose weight ``[Cin, Cout/groups, k...]`` followed by BN: output channel of ``W[i, j]`` is
    ``(i // (Cin/groups)) * (Cout/groups) + j``."""
    w = np.asarray(weight, np.float64)
    cin, cout_g = w.shape[0], w.shape[1]
    if cin % groups:
        raise ValueError(f"Cin={cin} is not divisible by groups={groups}")
    if cout_g * groups != bn.channels:
        raise ValueError(f"ConvTranspose has {cout_g * groups} output channels, BatchNorm has {bn.channels}")
    out_idx = (np.arange(cin)[:, None] // (cin // groups)) * cout_g + np.arange(cout_g)[None, :]
    s = bn.scale[out_idx].reshape((cin, cout_g) + (1,) * (w.ndim - 2))
    b = np.zeros(bn.channels, np.float64) if bias is None else np.asarray(bias, np.float64).reshape(-1)
    return _finish(w * s, (b - bn.mean) * bn.scale + bn.beta, dtype)


# ----------------------------------------------------------------------------------------------- ONNX

@dataclass(frozen=True)
class OnnxNode:
    """A graph node: name, op type, input / output tensor names and decoded attributes."""

    name: str
    op_type: str
    inputs: Tuple[str, ...]
    outputs: Tuple[str, ...]
    attrs: Dict[str, Any] = field(default_factory=dict)
    index: int = 0


@dataclass(frozen=True)
class ConvParams:
    """Weights and attributes of a Conv / ConvTranspose node. ``pads`` is the explicit attribute (zeros when
    absent); with ``auto_pad`` ``SAME_UPPER`` / ``SAME_LOWER`` the real padding depends on the input size and must
    be derived from it (PyTorch exports always write explicit pads and ``NOTSET``)."""

    weight: np.ndarray
    bias: Optional[np.ndarray]
    strides: Tuple[int, ...]
    pads: Tuple[int, ...]
    dilations: Tuple[int, ...]
    group: int
    kernel_shape: Tuple[int, ...]
    output_padding: Tuple[int, ...] = ()
    auto_pad: str = "NOTSET"
    output_shape: Tuple[int, ...] = ()


@dataclass(frozen=True)
class GemmParams:
    """Weights and attributes of a Gemm node (``weight`` exactly as stored; see ``trans_b``)."""

    weight: np.ndarray
    bias: Optional[np.ndarray]
    trans_a: bool
    trans_b: bool
    alpha: float
    beta: float


def _attr_value(attr) -> Any:
    from onnx import helper

    value = helper.get_attribute_value(attr)
    if isinstance(value, bytes):
        return value.decode("utf-8", "replace")
    if isinstance(value, list) and value and isinstance(value[0], bytes):
        return [v.decode("utf-8", "replace") for v in value]
    return value


class OnnxWeights:
    """Initializers and Constant tensors of an ONNX file, addressable by tensor name or by consuming node.

    Example::

        w = OnnxWeights("pts_backbone_neck_head_centerpoint.onnx")
        conv = w.conv("/backbone/blocks.0/blocks.0.0/Conv")            # ConvParams(weight, bias, strides, ...)
        wf, bf = w.fold_conv_bn("/backbone/blocks.0/blocks.0.0/Conv")  # conv + its BN consumer, fp64 fold
        k = w.param("/dit/blocks.0/attn/MatMul", 1)                    # anonymous MatMul weight by slot
        for node in w.nodes("Conv", "/backbone/*"):                    # graph order, glob on node names
            print(node.name, w.conv(node.name).weight.shape)
    """

    def __init__(self, path: PathLike, *, load_external_data: bool = True):
        import onnx

        self.path = Path(path)
        model = onnx.load(str(self.path), load_external_data=False)
        if load_external_data:
            self._load_external(model)
        self._sha256: Optional[str] = None
        graph = model.graph
        self._tensors: Dict[str, Any] = {t.name: t for t in graph.initializer}
        self._cache: Dict[str, np.ndarray] = {}
        self._nodes: Dict[str, OnnxNode] = {}
        self._order: List[str] = []
        self._producer: Dict[str, str] = {}
        self._consumers: Dict[str, List[str]] = {}
        self._alias: Dict[str, str] = {}
        for i, n in enumerate(graph.node):
            name = n.name or f"{n.op_type}_{i}"
            if name in self._nodes:
                name = f"{name}#{i}"
            node = OnnxNode(name, n.op_type, tuple(n.input), tuple(n.output),
                            {a.name: _attr_value(a) for a in n.attribute}, i)
            self._nodes[name] = node
            self._order.append(name)
            for out in n.output:
                self._producer[out] = name
            for inp in n.input:
                if inp:
                    self._consumers.setdefault(inp, []).append(name)
            if n.op_type == "Constant" and n.output:
                self._add_constant(n.output[0], n)
        for name in self._order:  # Identity of a constant tensor is that tensor
            node = self._nodes[name]
            if node.op_type == "Identity" and node.inputs and node.outputs:
                self._alias[node.outputs[0]] = node.inputs[0]
        self.input_names = tuple(i.name for i in graph.input if i.name not in self._tensors)
        self.output_names = tuple(o.name for o in graph.output)
        self.opset = {o.domain or "ai.onnx": int(o.version) for o in model.opset_import}

    def _load_external(self, model) -> None:
        from onnx import TensorProto
        from onnx import external_data_helper as ext

        base = self.path.parent.resolve()
        for tensor in model.graph.initializer:
            if tensor.data_location != TensorProto.EXTERNAL:
                continue
            location = {e.key: e.value for e in tensor.external_data}.get("location", "")
            target = (base / location).resolve()
            if base not in target.parents:
                raise ValueError(f"external data of {tensor.name!r} points outside {base}: {location!r}")
            ext.load_external_data_for_tensor(tensor, str(base))
            tensor.data_location = TensorProto.DEFAULT
            del tensor.external_data[:]

    def _add_constant(self, output: str, node) -> None:
        from onnx import numpy_helper

        for attr in node.attribute:
            if attr.name == "value":
                self._tensors[output] = attr.t
            elif attr.name in ("value_float", "value_floats", "value_int", "value_ints"):
                dt = np.float32 if "float" in attr.name else np.int64
                self._cache[output] = np.asarray(_attr_value(attr), dtype=dt)
                self._tensors[output] = numpy_helper.from_array(self._cache[output], output)

    # ---- tensors ---------------------------------------------------------------------------------
    @property
    def sha256(self) -> str:
        """sha256 of the ``.onnx`` file (computed on first use; a cache-key ingredient)."""
        if self._sha256 is None:
            self._sha256 = file_sha256(self.path)
        return self._sha256

    def _resolve(self, name: str) -> str:
        seen = set()
        while name in self._alias and name not in self._tensors:
            if name in seen:
                raise ValueError(f"Identity cycle at {name!r}")
            seen.add(name)
            name = self._alias[name]
        return name

    def has(self, name: str) -> bool:
        """True if ``name`` is an initializer / Constant output (or an Identity of one)."""
        return self._resolve(name) in self._tensors

    def array(self, name: str) -> np.ndarray:
        """The constant tensor ``name`` as numpy (cached; treat it as read-only)."""
        from onnx import numpy_helper

        key = self._resolve(name)
        if key not in self._tensors:
            raise KeyError(f"{name!r} is not an initializer or Constant output of {self.path.name}")
        if key not in self._cache:
            self._cache[key] = numpy_helper.to_array(self._tensors[key])
        return self._cache[key]

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

    def state_dict(self) -> Dict[str, np.ndarray]:
        """Every constant tensor by name (the export's names; prefer node addressing for stability)."""
        return {name: self.array(name) for name in self._tensors}

    # ---- nodes -----------------------------------------------------------------------------------
    def node(self, name: str) -> OnnxNode:
        try:
            return self._nodes[name]
        except KeyError:
            close = [n for n in self._order if name in n][:5]
            raise KeyError(f"no node {name!r} in {self.path.name}" + (f"; similar: {close}" if close else "")) from None

    def nodes(self, op_type: Optional[str] = None, pattern: Optional[str] = None, *,
              regex: bool = False) -> List[OnnxNode]:
        """Nodes in graph order, filtered by op type and a glob (or ``regex=True``) on the node name."""
        rx = re.compile(pattern) if (pattern and regex) else None
        out = []
        for name in self._order:
            node = self._nodes[name]
            if op_type is not None and node.op_type != op_type:
                continue
            if pattern is not None and not (rx.search(name) if rx else fnmatch.fnmatchcase(name, pattern)):
                continue
            out.append(node)
        return out

    def find_node(self, pattern: str, op_type: Optional[str] = None, *, regex: bool = False) -> OnnxNode:
        """Exactly one node matching ``pattern`` (and ``op_type``), else ``KeyError``."""
        found = self.nodes(op_type, pattern, regex=regex)
        if len(found) != 1:
            raise KeyError(f"{len(found)} nodes match {pattern!r} (op_type={op_type}): {[n.name for n in found[:8]]}")
        return found[0]

    def producer(self, tensor: str) -> Optional[OnnxNode]:
        name = self._producer.get(tensor)
        return self._nodes[name] if name else None

    def consumers(self, tensor: str) -> List[OnnxNode]:
        """Nodes reading tensor ``tensor`` (directly; Identity nodes are listed as consumers)."""
        return [self._nodes[n] for n in self._consumers.get(tensor, [])]

    def consumer_of(self, tensor: str, op_type: Optional[str] = None) -> OnnxNode:
        """The single consumer of ``tensor`` (optionally of ``op_type``), else ``KeyError``."""
        found = [n for n in self.consumers(tensor) if op_type is None or n.op_type == op_type]
        if len(found) != 1:
            raise KeyError(f"{tensor!r} has {len(found)} consumers of type {op_type}: {[n.name for n in found]}")
        return found[0]

    def param(self, node_name: str, index: int) -> np.ndarray:
        """The constant feeding input slot ``index`` of node ``node_name``."""
        node = self.node(node_name)
        if index >= len(node.inputs) or not node.inputs[index]:
            raise KeyError(f"node {node_name!r} ({node.op_type}) has no input {index}")
        name = node.inputs[index]
        if not self.has(name):
            raise KeyError(f"input {index} of {node_name!r} is {name!r}, a runtime tensor, not a constant")
        return self.array(name)

    def params(self, node_name: str) -> Dict[int, np.ndarray]:
        """``{slot: array}`` for every constant input of the node."""
        node = self.node(node_name)
        return {i: self.array(t) for i, t in enumerate(node.inputs) if t and self.has(t)}

    def conv(self, node_name: str) -> ConvParams:
        """Conv / ConvTranspose weights (+ bias) and attributes with ONNX defaults filled in."""
        node = self.node(node_name)
        if node.op_type not in ("Conv", "ConvTranspose"):
            raise ValueError(f"{node_name!r} is {node.op_type}, not Conv/ConvTranspose")
        w = self.param(node_name, 1)
        b = self.param(node_name, 2) if len(node.inputs) > 2 and node.inputs[2] else None
        k = tuple(node.attrs.get("kernel_shape", w.shape[2:]))
        nd = len(k)
        attrs = node.attrs
        return ConvParams(w, b, tuple(attrs.get("strides", (1,) * nd)), tuple(attrs.get("pads", (0,) * 2 * nd)),
                          tuple(attrs.get("dilations", (1,) * nd)), int(attrs.get("group", 1)), k,
                          tuple(attrs.get("output_padding", ())), str(attrs.get("auto_pad", "NOTSET")),
                          tuple(attrs.get("output_shape", ())))

    def gemm(self, node_name: str) -> GemmParams:
        node = self.node(node_name)
        if node.op_type != "Gemm":
            raise ValueError(f"{node_name!r} is {node.op_type}, not Gemm")
        b = self.param(node_name, 2) if len(node.inputs) > 2 and node.inputs[2] else None
        return GemmParams(self.param(node_name, 1), b, bool(node.attrs.get("transA", 0)),
                          bool(node.attrs.get("transB", 0)), float(node.attrs.get("alpha", 1.0)),
                          float(node.attrs.get("beta", 1.0)))

    def matmul_weight(self, node_name: str) -> np.ndarray:
        """The constant operand of a MatMul (``[in, out]`` when it is the right operand)."""
        node = self.node(node_name)
        consts = [i for i, t in enumerate(node.inputs[:2]) if t and self.has(t)]
        if len(consts) != 1:
            raise ValueError(f"MatMul {node_name!r} has {len(consts)} constant operands")
        return self.array(node.inputs[consts[0]])

    def batchnorm(self, node_name: str) -> BatchNorm:
        node = self.node(node_name)
        if node.op_type != "BatchNormalization":
            raise ValueError(f"{node_name!r} is {node.op_type}, not BatchNormalization")
        return BatchNorm(self.param(node_name, 1), self.param(node_name, 2), self.param(node_name, 3),
                         self.param(node_name, 4), float(node.attrs.get("epsilon", 1e-5)))

    def fold_conv_bn(self, conv_node: str, bn_node: Optional[str] = None, *,
                     dtype: Optional[Any] = np.float32) -> Tuple[np.ndarray, np.ndarray]:
        """Conv / ConvTranspose weights with the following BatchNormalization folded in (float64, one rounding).
        ``bn_node`` defaults to the single BatchNormalization consuming the conv output."""
        conv = self.conv(conv_node)
        if bn_node is None:
            bn_node = self.consumer_of(self.node(conv_node).outputs[0], "BatchNormalization").name
        bn = self.batchnorm(bn_node)
        if self.node(conv_node).op_type == "ConvTranspose":
            return fold_bn_conv_transpose(conv.weight, conv.bias, bn, groups=conv.group, dtype=dtype)
        return fold_bn_conv(conv.weight, conv.bias, bn, dtype=dtype)

    def __repr__(self) -> str:
        return f"OnnxWeights({self.path.name}: {len(self._nodes)} nodes, {len(self._tensors)} constants)"


# ------------------------------------------------------------------------------------- state dicts

def _torch_to_numpy(t: Any) -> np.ndarray:
    t = t.detach().cpu()
    if str(t.dtype) == "torch.bfloat16":
        t = t.float()
    return t.numpy()


def load_safetensors(path: PathLike) -> Dict[str, np.ndarray]:
    """All tensors of a ``.safetensors`` file as numpy (bf16 tensors are widened to float32 through torch)."""
    from safetensors import safe_open

    out: Dict[str, np.ndarray] = {}
    try:
        with safe_open(str(path), framework="numpy") as f:
            for k in f.keys():
                out[k] = f.get_tensor(k)
        return out
    except (TypeError, ValueError):  # bf16 has no numpy dtype
        out.clear()
    with safe_open(str(path), framework="pt") as f:
        for k in f.keys():
            out[k] = _torch_to_numpy(f.get_tensor(k))
    return out


_CHECKPOINT_KEYS = ("state_dict", "model", "model_state_dict", "module", "net", "params")


def load_torch_checkpoint(path: PathLike, *, key: Optional[str] = None) -> Dict[str, np.ndarray]:
    """A ``.pth`` / ``.pt`` / ``.ckpt`` checkpoint loaded with ``torch.load(weights_only=True)`` (data only: no
    pickled code runs) -> ``{name: numpy}``. ``key`` selects a nested dict; otherwise the first of
    ``state_dict`` / ``model`` / ``model_state_dict`` / ``module`` / ``net`` / ``params`` that holds tensors."""
    import torch

    obj = torch.load(str(path), map_location="cpu", weights_only=True)
    if key is not None:
        obj = obj[key]
    elif isinstance(obj, Mapping) and not any(hasattr(v, "detach") for v in obj.values()):
        for k in _CHECKPOINT_KEYS:
            if isinstance(obj.get(k), Mapping):
                obj = obj[k]
                break
    if not isinstance(obj, Mapping):
        raise ValueError(f"{path}: no tensor dict found (pass key=)")
    return {str(k): _torch_to_numpy(v) for k, v in obj.items() if hasattr(v, "detach")}


def load_state_dict(path: PathLike, **kwargs: Any) -> "StateDict":
    """``.safetensors`` / ``.pth`` / ``.pt`` / ``.ckpt`` / ``.npz`` -> :class:`StateDict` (by file suffix)."""
    p = Path(path)
    suffix = p.suffix.lower()
    if suffix == ".safetensors":
        tensors = load_safetensors(p)
    elif suffix in (".pth", ".pt", ".ckpt", ".bin"):
        tensors = load_torch_checkpoint(p, **kwargs)
    elif suffix == ".npz":
        with np.load(p, allow_pickle=False) as z:
            tensors = {k: z[k] for k in z.files}
    else:
        raise ValueError(f"unsupported weights file {p.name}")
    return StateDict(tensors)


class StateDict(Mapping):
    """A read-only ``{name: numpy}`` view with prefix scoping: ``sd.sub("backbone.layer1.")["0.conv1.weight"]``."""

    def __init__(self, tensors: Mapping[str, np.ndarray], prefix: str = ""):
        self._tensors = tensors
        self.prefix = prefix

    def __getitem__(self, name: str) -> np.ndarray:
        try:
            return self._tensors[self.prefix + name]
        except KeyError:
            raise KeyError(f"{self.prefix + name!r} not in the state dict") from None

    def __iter__(self) -> Iterator[str]:
        n = len(self.prefix)
        return (k[n:] for k in self._tensors if k.startswith(self.prefix))

    def __len__(self) -> int:
        return sum(1 for _ in self)

    def sub(self, prefix: str) -> "StateDict":
        return StateDict(self._tensors, self.prefix + prefix)

    def strip(self, prefix: str) -> "StateDict":
        """Drop a leading ``prefix`` (e.g. ``module.``) from every key that has it."""
        return StateDict({(k[len(prefix):] if k.startswith(prefix) else k): v for k, v in self._tensors.items()},
                         self.prefix)

    def bn(self, name: str, eps: float = 1e-5) -> BatchNorm:
        """BatchNorm ``<name>.{weight,bias,running_mean,running_var}``."""
        return BatchNorm.from_state_dict(self, name, eps)


# ----------------------------------------------------------------------------------------- device cache

def file_sha256(path: PathLike, chunk: int = 1 << 20) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for block in iter(lambda: f.read(chunk), b""):
            h.update(block)
    return h.hexdigest()


def cache_root() -> Path:
    """``$TT_CACHE_PATH`` (tt-model containers mount ``/tensor-cache`` there), else ``~/.cache/ttaw``."""
    env = os.environ.get("TT_CACHE_PATH")
    return Path(env) if env else Path.home() / ".cache" / "ttaw"


_SAFE = re.compile(r"[^A-Za-z0-9_.-]+")


class WeightCache:
    """Prepared device tensors cached on disk as ttnn ``.tensorbin`` files.

    ``namespace`` names the model (``"tt_centerpoint/base"``); ``source_digest`` identifies the source weights
    (e.g. ``OnnxWeights.sha256``) so a new revision gets a fresh directory. ``version`` names the preparation
    recipe and is part of the directory too: the cache outlives images (tt-model mounts the host's
    ``~/.cache/tt-model/<name>/tensors`` at ``/tensor-cache``), so an entry built by older preparation code from the
    same source weights would otherwise be served to newer code; bump it (e.g. the bundle ``__version__`` plus a
    prep revision) whenever the code that builds the tensors changes. Each entry is keyed by name, dtype and layout;
    a loaded entry whose dtype, layout or (with ``shape=``) logical shape differs from the request is rebuilt.
    ``enabled=None`` reads ``TTAW_WEIGHT_CACHE`` once (``0`` disables: every call rebuilds)."""

    def __init__(self, namespace: str, source_digest: str = "", *, version: str = "",
                 root: Optional[PathLike] = None, enabled: Optional[bool] = None):
        if enabled is None:
            enabled = os.environ.get("TTAW_WEIGHT_CACHE", "1").strip() not in ("0", "false", "off", "no")
        self.enabled = bool(enabled)
        parts = [p for p in namespace.strip("/").split("/") if p]
        if not parts or any(p in (".", "..") for p in parts):
            raise ValueError(f"bad cache namespace {namespace!r}")
        leaf = (source_digest or "nodigest")[:16] + (f"-{_SAFE.sub('_', str(version))}" if version else "")
        self.version = str(version)
        self.dir = Path(root) if root is not None else cache_root()
        self.dir = self.dir.joinpath(*(_SAFE.sub("_", p) for p in parts), leaf)
        self.hits = 0
        self.misses = 0

    def path(self, key: str, dtype: Any, layout: Any) -> Path:
        tag = f"{key}|{dtype}|{layout}"
        short = hashlib.sha1(tag.encode()).hexdigest()[:8]
        dname = getattr(dtype, "name", str(dtype)).lower()
        lname = getattr(layout, "name", str(layout)).lower()
        return self.dir / f"{_SAFE.sub('_', key)[:120]}.{dname}.{lname}.{short}.tensorbin"

    def get(self, key: str, make: Callable[[], Any], *, dtype: Any, layout: Any = None, device: Any = None,
            memory_config: Any = None, shape: Optional[Tuple[int, ...]] = None):
        """The cached host tensor for ``key`` (built with ``make()`` -> numpy / torch on a miss), moved to
        ``device`` with ``memory_config`` when given. ``shape``: the expected logical shape (an entry of another
        shape is rebuilt)."""
        import ttnn

        from .tensors import to_host_tensor, ttnn_dtype

        dtype = ttnn_dtype(dtype)
        layout = ttnn.TILE_LAYOUT if layout is None else layout
        host = None
        p = self.path(key, dtype, layout)
        if self.enabled and p.is_file():
            try:
                host = ttnn.load_tensor(p)
            except Exception:  # noqa: BLE001 -- unreadable / old format: rebuild below
                host = None
            if host is not None and (host.dtype != dtype or host.layout != layout
                                     or (shape is not None and tuple(host.shape) != tuple(shape))):
                host = None   # an entry of another recipe: rebuild
            if host is not None:
                self.hits += 1
        if host is None:
            self.misses += 1
            host = to_host_tensor(make(), dtype, layout, shape=shape)
            if self.enabled:
                p.parent.mkdir(parents=True, exist_ok=True)
                tmp = p.with_name(f".{p.stem}.{os.getpid()}.tmp.tensorbin")
                ttnn.dump_tensor(tmp, host)
                os.replace(tmp, p)
        if device is None:
            return host
        return ttnn.to_device(host, device, memory_config=memory_config or ttnn.DRAM_MEMORY_CONFIG)

    def clear(self) -> int:
        """Delete this namespace/digest directory's entries; returns how many files were removed."""
        n = 0
        if self.dir.is_dir():
            for f in self.dir.glob("*.tensorbin"):
                f.unlink()
                n += 1
        return n