"""Architecture selection including the projected scalar warmup experiment."""
import torch

from tinyasr.codec_factory_v1 import load_codec_checkpoint as load_previous
from tinyasr.codec_projected_warmup_v1 import ARCHITECTURE, ProjectedWarmupCodec


def load_codec_checkpoint(path):
    checkpoint = torch.load(path, map_location='cpu', weights_only=False)
    if checkpoint.get('architecture') != ARCHITECTURE:
        return load_previous(path)
    model = ProjectedWarmupCodec()
    model.load_state_dict(checkpoint['codec'], strict=True)
    return model.eval().requires_grad_(False)