"""Model registry — maps config name to model class."""

from __future__ import annotations

from typing import Any

import torch.nn as nn

MODEL_REGISTRY: dict[str, type[nn.Module]] = {}


def register_model(name: str):
    """Decorator to register a model class."""

    def decorator(cls):
        MODEL_REGISTRY[name] = cls
        return cls

    return decorator


def get_model(name: str, **kwargs: Any) -> nn.Module:
    """Instantiate a registered model by name."""
    if name not in MODEL_REGISTRY:
        raise ValueError(f"Unknown model '{name}'. Available: {list(MODEL_REGISTRY.keys())}")
    return MODEL_REGISTRY[name](**kwargs)


# Auto-import submodules to trigger registration
from softasp_solver.sudoku import baselines as _sudoku_baselines  # noqa: F401, E402
from softasp_solver.sudoku import softasp_sudoku as _sudoku  # noqa: F401, E402