"""Storage dtype of every exported tensor.

BF16 for every weight that BF16 autocast consumes in BF16 anyway (matrices,
convolutions, biases, and norm gains and biases applied to BF16 activations):
storing them in BF16 gives bit-identical outputs. FP32 where the value takes
part in FP32 arithmetic: the frozen DINOv3 residual stream (CLS and register
tokens, layer scales, LayerNorms), its RoPE periods, the input and
layer-standardization statistics, the latent statistics used for whitening,
and the upsamplers' 3x3 weights and fold matrices (the 4x4 transposed kernel is
folded from them outside autocast, in their storage dtype).
"""

from __future__ import annotations

import re
from typing import TYPE_CHECKING

import torch

if TYPE_CHECKING:
    from collections.abc import Mapping

    from torch import Tensor

_FP32_NAMES = frozenset(
    {
        "latent_mean",
        "latent_var",
        "encoder.pixel_mean",
        "encoder.pixel_std",
        "encoder.layer_mean",
        "encoder.layer_std",
        "encoder.backbone.cls_token",
        "encoder.backbone.reg_token",
        "encoder.backbone.rope.periods",
        "encoder.backbone.norm.weight",
        "encoder.backbone.norm.bias",
    }
)
_UPSAMPLER = re.compile(
    r"decoder\.conv_up_head\.(handoff_upsample|upsamples\.\d+)\.(conv\.weight|fold)"
)
_BACKBONE_BLOCK_PREFIX = "encoder.backbone.blocks."
_FP32_BLOCK_SUFFIXES = (
    ".gamma_1",
    ".gamma_2",
    ".norm1.weight",
    ".norm1.bias",
    ".norm2.weight",
    ".norm2.bias",
)


def storage_dtype(name: str) -> torch.dtype:
    """BF16 or FP32 storage of one exported tensor, by its state-dict name."""

    fp32 = (
        name in _FP32_NAMES
        or (
            name.startswith(_BACKBONE_BLOCK_PREFIX)
            and name.endswith(_FP32_BLOCK_SUFFIXES)
        )
        or _UPSAMPLER.fullmatch(name) is not None
    )
    return torch.float32 if fp32 else torch.bfloat16


def require_storage_policy(state: Mapping[str, Tensor]) -> None:
    """Require every tensor of an artifact to be stored in its policy dtype.

    Raises:
        ValueError: Naming every tensor stored in another dtype.
    """

    wrong = sorted(
        name for name, tensor in state.items() if tensor.dtype != storage_dtype(name)
    )
    if wrong:
        raise ValueError(f"Tensors not stored in their policy dtype: {wrong}")