src/models/teacher_strict.py

"""Strict, auditable L2CS checkpoint loading.

This module intentionally does not replace the historical loader in ``teacher.py``.
It is the loader for new, versioned KD artifacts only.
"""
from __future__ import annotations

from dataclasses import asdict, dataclass
import hashlib
from pathlib import Path
from typing import Mapping

import torch

from src.models.teacher import L2CS


PREFIX_MAP = {
    "conv1.": "backbone.0.",
    "bn1.": "backbone.1.",
    "layer1.": "backbone.4.",
    "layer2.": "backbone.5.",
    "layer3.": "backbone.6.",
    "layer4.": "backbone.7.",
}


@dataclass(frozen=True)
class StrictLoadAudit:
    checkpoint_path: str
    checkpoint_sha256: str
    checkpoint_tensor_count: int
    mapped_tensor_count: int
    model_tensor_count: int
    allowed_model_only_keys: tuple[str, ...]
    inference_sha256: str | None = None


def file_sha256(path: str | Path) -> str:
    digest = hashlib.sha256()
    with Path(path).open("rb") as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest().upper()


def _extract_state_dict(checkpoint: object) -> Mapping[str, torch.Tensor]:
    if not isinstance(checkpoint, Mapping):
        raise TypeError(f"Checkpoint must be a mapping, received {type(checkpoint)!r}")
    state = checkpoint.get("state_dict", checkpoint)
    if not isinstance(state, Mapping):
        raise TypeError("Checkpoint 'state_dict' must be a mapping")
    tensors = {str(key): value for key, value in state.items() if torch.is_tensor(value)}
    if not tensors:
        raise ValueError("Checkpoint contains no tensor parameters")
    return tensors


def remap_l2cs_resnet_key(key: str) -> str:
    """Map raw torchvision/L2CS ResNet names into the wrapper namespace."""
    if key.startswith("module."):
        key = key[len("module.") :]
    for source, target in PREFIX_MAP.items():
        if key.startswith(source):
            return target + key[len(source) :]
    return key


def _remap_and_validate(
    state: Mapping[str, torch.Tensor], model: L2CS
) -> dict[str, torch.Tensor]:
    model_state = model.state_dict()
    remapped: dict[str, torch.Tensor] = {}
    source_for_target: dict[str, str] = {}
    problems: list[str] = []

    for source_key, tensor in state.items():
        target_key = remap_l2cs_resnet_key(source_key)
        if target_key in remapped:
            problems.append(
                f"duplicate mapping: {source_key!r} and {source_for_target[target_key]!r} "
                f"both map to {target_key!r}"
            )
            continue
        if target_key not in model_state:
            problems.append(f"unexpected checkpoint tensor: {source_key!r} -> {target_key!r}")
            continue
        expected = model_state[target_key]
        if tuple(tensor.shape) != tuple(expected.shape):
            problems.append(
                f"shape mismatch for {source_key!r} -> {target_key!r}: "
                f"checkpoint={tuple(tensor.shape)}, model={tuple(expected.shape)}"
            )
            continue
        if tensor.dtype != expected.dtype:
            problems.append(
                f"dtype mismatch for {source_key!r} -> {target_key!r}: "
                f"checkpoint={tensor.dtype}, model={expected.dtype}"
            )
            continue
        remapped[target_key] = tensor
        source_for_target[target_key] = source_key

    model_only = sorted(set(model_state) - set(remapped))
    if model_only != ["idx_tensor"]:
        problems.append(f"missing model tensors other than allowed idx_tensor: {model_only}")
    if len(remapped) != len(state):
        problems.append(f"mapped {len(remapped)} of {len(state)} checkpoint tensors")
    if problems:
        raise RuntimeError("Strict L2CS load validation failed:\n- " + "\n- ".join(problems))
    return remapped


def fixed_inference_sha256(model: L2CS, device: torch.device) -> str:
    """Hash deterministic logits at the recovered checkpoint's documented 448px input."""
    values = torch.linspace(-1.0, 1.0, steps=3 * 448 * 448, dtype=torch.float32)
    fixed_input = values.reshape(1, 3, 448, 448).to(device)
    with torch.inference_mode():
        pitch, yaw = model(fixed_input)
    output = torch.cat((pitch, yaw), dim=1).detach().cpu().contiguous().numpy()
    return hashlib.sha256(output.tobytes(order="C")).hexdigest().upper()


def load_teacher_model_strict(
    checkpoint_path: str | Path,
    *,
    backbone: str = "resnet50",
    num_bins: int = 90,
    device: str | torch.device = "cpu",
    compute_inference_hash: bool = True,
) -> tuple[L2CS, StrictLoadAudit]:
    """Load every trained checkpoint tensor or fail before inference."""
    checkpoint_path = Path(checkpoint_path).resolve()
    target_device = torch.device(device)
    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
    state = _extract_state_dict(checkpoint)
    model = L2CS(backbone_name=backbone, num_bins=num_bins, pretrained=False)
    remapped = _remap_and_validate(state, model)

    # idx_tensor is generated deterministically by the wrapper; all trained tensors were
    # already checked above. No incompatibility result is ignored.
    completed = dict(remapped)
    completed["idx_tensor"] = model.state_dict()["idx_tensor"]
    model.load_state_dict(completed, strict=True)
    model.to(target_device)
    model.eval()

    inference_hash = fixed_inference_sha256(model, target_device) if compute_inference_hash else None
    audit = StrictLoadAudit(
        checkpoint_path=str(checkpoint_path),
        checkpoint_sha256=file_sha256(checkpoint_path),
        checkpoint_tensor_count=len(state),
        mapped_tensor_count=len(remapped),
        model_tensor_count=len(model.state_dict()),
        allowed_model_only_keys=("idx_tensor",),
        inference_sha256=inference_hash,
    )
    return model, audit


def audit_to_dict(audit: StrictLoadAudit) -> dict:
    return asdict(audit)