"""Model registry — name → builder.

Importing the architecture modules has side effects (registers each
class with REGISTRY) so callers can do::

    from training.models import get_model
    cls = get_model("cnn")

without knowing which file defines it.
"""
from __future__ import annotations

from typing import Any, Callable

REGISTRY: dict[str, Callable[..., "BaseModel"]] = {}


def register(name: str):
    def decorator(cls):
        if name in REGISTRY:
            raise ValueError(f"model {name!r} already registered")
        REGISTRY[name] = cls
        cls.__model_name__ = name
        return cls
    return decorator


def get_model(name: str):
    if name not in REGISTRY:
        raise KeyError(
            f"model {name!r} not registered; known: {sorted(REGISTRY)}"
        )
    return REGISTRY[name]


# Eager-import the implementations so the registry is populated.
# Order matters only for which "kind" gets imported first — all are listed.
from training.models import gbt              # noqa: F401,E402
from training.models import mlp              # noqa: F401,E402
from training.models import cnn              # noqa: F401,E402
from training.models import gru              # noqa: F401,E402
from training.models import lstm             # noqa: F401,E402
from training.models import transformer      # noqa: F401,E402
from training.models import transformer_ssl  # noqa: F401,E402

from training.models._base import BaseModel  # noqa: E402,F401