"""
模型定义 - 使用timm预训练模型，只微调分类头
"""

import torch
import torch.nn as nn
import timm

import config


def create_model(pretrained: bool = True, freeze_backbone: bool = True) -> nn.Module:
    """
    创建模型，只微调分类头

    Args:
        pretrained: 是否使用预训练权重
        freeze_backbone: 是否冻结backbone

    Returns:
        PyTorch模型
    """
    # 使用timm创建模型
    model = timm.create_model(
        config.MODEL_NAME, pretrained=pretrained, num_classes=config.NUM_CLASSES
    )

    if freeze_backbone:
        # 冻结除分类头外的所有层
        for name, param in model.named_parameters():
            if "fc" not in name and "classifier" not in name and "head" not in name:
                param.requires_grad = False

        # 统计可训练参数
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        total_params = sum(p.numel() for p in model.parameters())

        print(f"模型: {config.MODEL_NAME}")
        print(f"总参数量: {total_params:,}")
        print(f"可训练参数量: {trainable_params:,}")
        print(f"冻结参数量: {total_params - trainable_params:,}")

    return model


def load_model(model_path: str, device: str = None) -> nn.Module:
    """
    加载训练好的模型

    Args:
        model_path: 模型权重路径
        device: 运行设备

    Returns:
        加载权重后的模型
    """
    if device is None:
        device = config.DEVICE

    model = create_model(pretrained=False, freeze_backbone=False)

    # 加载权重
    checkpoint = torch.load(model_path, map_location=device)
    if "model_state_dict" in checkpoint:
        model.load_state_dict(checkpoint["model_state_dict"])
    else:
        model.load_state_dict(checkpoint)

    model = model.to(device)
    model.eval()

    return model


if __name__ == "__main__":
    # 测试模型创建
    model = create_model(pretrained=True, freeze_backbone=True)

    # 测试前向传播
    x = torch.randn(1, 3, config.IMAGE_SIZE, config.IMAGE_SIZE)
    with torch.no_grad():
        output = model(x)

    print(f"\n输入形状: {x.shape}")
    print(f"输出形状: {output.shape}")