"""timm によるモデル生成（SPEC_02「モデル」）.

`model.name` に timm のモデル名を渡すだけでアーキテクチャを差し替えられるようにする．
"""

from __future__ import annotations

import timm
from torch import nn

from tias_judge_model.data.classes import NUM_CLASSES


def create_model(name: str, *, pretrained: bool, num_classes: int = NUM_CLASSES) -> nn.Module:
    """timm のモデルを出力 `num_classes` クラスで作る.

    Args:
        name: timm のモデル名（例: `convnext_tiny`・`efficientnet_b0`・`vit_small_patch16_224`）．
        pretrained: ImageNet 事前学習の重みを使うか（初回はダウンロードする）．
        num_classes: 出力クラス数．

    Returns:
        入力 `(B, 3, H, W)` の float テンソルから `(B, num_classes)` のロジットを返すモデル．

    Raises:
        ValueError: timm に無いモデル名の場合．
    """
    if not timm.is_model(name):
        raise ValueError(f"timm に無いモデル名である: {name!r}")
    return timm.create_model(name, pretrained=pretrained, num_classes=num_classes)


def count_parameters(model: nn.Module) -> int:
    """モデルのパラメータ数（学習対象外も含む）を返す.

    Args:
        model: モデル．

    Returns:
        パラメータの要素数の合計．
    """
    return sum(p.numel() for p in model.parameters())