"""Strict loading with explicit architecture selection for codec experiments."""
import torch

from tinyasr.codec_normalized_v1 import ARCHITECTURE, UnitScaleCodec
from tinyasr.tinycall import TinyCallCodec


def load_codec_checkpoint(path):
    checkpoint = torch.load(path, map_location='cpu', weights_only=False)
    architecture = checkpoint.get('architecture', 'tinycall_independent_v1')
    if architecture == ARCHITECTURE:
        model = UnitScaleCodec()
    elif architecture == 'tinycall_independent_v1':
        model = TinyCallCodec()
    else:
        raise ValueError('Unsupported codec checkpoint architecture')
    model.load_state_dict(checkpoint['codec'], strict=True)
    return model.eval().requires_grad_(False)