import torch
import torch.nn as nn

class PatchEmbedding(nn.Module):
    """이미지를 패치로 분할하고 임베딩합니다."""
    def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_dim=768):
        super().__init__()
        self.num_patches = (img_size // patch_size) ** 2  # 196
        # 패치 분할 + 선형 변환을 Conv2d로 구현
        self.proj = nn.Conv2d(in_channels, embed_dim,
                              kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        # x: [B, 3, 224, 224]
        x = self.proj(x)       # [B, 768, 14, 14]
        x = x.flatten(2)       # [B, 768, 196]
        x = x.transpose(1, 2)  # [B, 196, 768]
        return x

class ViTEmbedding(nn.Module):
    """CLS 토큰과 위치 임베딩을 추가합니다."""
    def __init__(self, img_size=224, patch_size=16, embed_dim=768):
        super().__init__()
        self.patch_embed = PatchEmbedding(img_size, patch_size, embed_dim=embed_dim)
        num_patches = self.patch_embed.num_patches

        # CLS 토큰: 이미지 전체를 대표하는 학습 가능 벡터
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        # 위치 임베딩: 패치의 공간적 위치 정보
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))

    def forward(self, x):
        B = x.shape[0]
        x = self.patch_embed(x)  # [B, 196, 768]

        # CLS 토큰 추가
        cls_tokens = self.cls_token.expand(B, -1, -1)  # [B, 1, 768]
        x = torch.cat([cls_tokens, x], dim=1)  # [B, 197, 768]

        # 위치 임베딩 추가
        x = x + self.pos_embed  # [B, 197, 768]
        return x

import timm

# ViT-Base/16 (16x16 패치, 12개 Transformer 블록)
model = timm.create_model('vit_base_patch16_224', pretrained=True, num_classes=10)

# 모델 구조 확인
print(f"패치 임베딩: {model.patch_embed}")
print(f"Transformer 블록 수: {len(model.blocks)}")
print(f"파라미터 수: {sum(p.numel() for p in model.parameters()):,}")

# 추론
dummy = torch.randn(1, 3, 224, 224)
output = model(dummy)
print(f"출력: {output.shape}")  # [1, 10]

def get_attention_map(model, image_tensor):
    """ViT의 마지막 레이어 Attention Map을 추출합니다."""
    model.eval()
    attention_maps = []

    # 어텐션 가중치를 저장하는 훅 등록
    def hook_fn(module, input, output):
        # output: (attn_output, attn_weights)
        attention_maps.append(output[1].detach())

    # 마지막 블록의 Attention에 훅 등록
    hook = model.blocks[-1].attn.register_forward_hook(hook_fn)

    with torch.no_grad():
        _ = model(image_tensor.unsqueeze(0))

    hook.remove()

    # CLS 토큰이 다른 패치에 주목하는 정도
    attn = attention_maps[0]  # [1, heads, 197, 197]
    attn = attn.mean(dim=1)   # 헤드 평균 [1, 197, 197]
    cls_attn = attn[0, 0, 1:]  # CLS → 패치 [196]
    cls_attn = cls_attn.reshape(14, 14)  # 2D 맵으로 변환
    return cls_attn.numpy()

# Swin Transformer 사용
swin = timm.create_model('swin_tiny_patch4_window7_224', pretrained=True, num_classes=10)

# 계층적 특징맵 추출 (탐지/세그멘테이션 백본으로 활용)
swin_backbone = timm.create_model(
    'swin_tiny_patch4_window7_224',
    pretrained=True,
    features_only=True,
)
features = swin_backbone(torch.randn(1, 3, 224, 224))
for i, f in enumerate(features):
    print(f"Stage {i}: {f.shape}")
# Stage 0: [1, 96, 56, 56]
# Stage 1: [1, 192, 28, 28]
# Stage 2: [1, 384, 14, 14]
# Stage 3: [1, 768, 7, 7]