"""Architecture dispatch shared by training and checkpoint inference."""

from vimeml.training.model import GPTConfig, TinyGPT
from vimeml.training.model_v2 import GPTV2Config, TinyGPTV2

ARCHITECTURES = {
    "tiny_gpt_v1": (GPTConfig, TinyGPT, "vimeml_tiny_gpt_v1"),
    "tiny_gpt_v2": (GPTV2Config, TinyGPTV2, "vimeml_tiny_gpt_v2"),
}


def specification(architecture):
    if architecture not in ARCHITECTURES:
        raise ValueError(f"Unsupported architecture: {architecture}")
    return ARCHITECTURES[architecture]


def configuration_for(architecture, values):
    config_type, _, _ = specification(architecture)
    return config_type(**values)


def create_model(architecture, config):
    config_type, model_type, _ = specification(architecture)
    if type(config) is not config_type:
        raise ValueError("Model configuration does not match architecture.")
    return model_type(config)


def checkpoint_format(architecture):
    return specification(architecture)[2]


def model_from_checkpoint(saved):
    architecture = next(
        (
            name
            for name, (_, _, format_name) in ARCHITECTURES.items()
            if format_name == saved.get("format")
        ),
        None,
    )
    if architecture is None:
        raise ValueError("Unsupported checkpoint format.")
    if (
        saved.get("architecture", architecture) != architecture
        or saved.get("config", {}).get("architecture", architecture) != architecture
    ):
        raise ValueError("Checkpoint architecture differs from its format.")
    model = create_model(architecture, configuration_for(architecture, saved["model_config"]))
    model.load_state_dict(saved["model"])
    return model