__all__ = ['build_model']

import importlib

from utils import Registry

META_ARCHITECTURES = Registry()
_MAMBA_MODELS = frozenset({"Mamba", "MambaA"})


def _load_optional_architecture(model_name):
    if model_name not in _MAMBA_MODELS or model_name in META_ARCHITECTURES:
        return

    try:
        importlib.import_module(".mamba", package=__package__)
    except ModuleNotFoundError as error:
        missing_package = error.name or "mamba_ssm"
        raise ModuleNotFoundError(
            f"Model {model_name!r} requires the optional Mamba backend "
            f"(missing package: {missing_package!r}). Install "
            "requirements-mamba.txt, or use MiniROAD/MiniROADA without it."
        ) from error


def build_model(cfg, device=None):
    model_name = cfg["model"]
    _load_optional_architecture(model_name)

    if model_name not in META_ARCHITECTURES:
        available = ", ".join(sorted(META_ARCHITECTURES))
        raise KeyError(
            f"Unknown model {model_name!r}. Available architectures: {available}"
        )

    model = META_ARCHITECTURES[model_name](cfg)
    return model.to(device)