import torch.nn as nn
import torch
import torch.optim as optim

class ViT(nn.Module):
    def __init__(self, image_size=28, patch_size=7, num_classes=10, dim=128, depth=6, heads=8, mlp_dim=256, dropout=0.1):
        super(ViT, self).__init__()

        assert image_size % patch_size == 0, 'image dimensions must be divisible by the patch size'
        num_patches = (image_size // patch_size) ** 2
        patch_dim = 1 * patch_size ** 2

        # 定义线性层将图像分块并映射到嵌入空间
        self.patch_embedding = nn.Linear(patch_dim, dim)

        # 位置编码
        # nn.Parameter是Pytorch中的一个类，用于将一个张量注册为模型的参数
        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches, dim))

        # Dropout层
        self.dropout = nn.Dropout(dropout)

        # Transformer编码器
        # 当 batch_first=True 时，输入和输出张量的形状为 (batch_size, seq_length, feature_dim)。当 batch_first=False 时，输入和输出张量的形状为 (seq_length, batch_size, feature_dim)。
        self.transformer = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(
                d_model=dim,
                nhead=heads,
                dim_feedforward=mlp_dim
                # batch_first=True
            ),
            num_layers=depth
        )

        # 分类头
        # nn.Identity()是一个空的层，它不执行任何操作，只是返回输入
        # self.to_cls_token = nn.Identity()
        # self.mlp_head = nn.Linear(dim, num_classes)
        self.mlp_head = nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, num_classes)
        )
    
    def forward(self, x):
        # x shape: [batch_size, 1, 28, 28]
        batch_size = x.size(0)
        x = x.view(batch_size, -1, 7*7)  # 将图像划分为7x7的Patch
        x = self.patch_embedding(x)  # [batch_size, num_patches, dim]
        x += self.pos_embedding  # 添加位置编码
        x = self.dropout(x)  # 应用Dropout

        x = x.permute(1, 0, 2)  # Transformer期望的输入形状：[seq_len, batch_size, embedding_dim]
        x = self.transformer(x)  # [序列长度, batch_size, dim]
        x = x.permute(1, 0, 2)  # 转回原来的形状：[batch_size, seq_len, dim]

        x = x.mean(dim=1)  # 对所有Patch取平均，x.mean(dim=1) 这一步是对所有 Patch 的特征向量取平均值，从而得到一个代表整个图像的全局特征向量。
        x = self.mlp_head(x)  # [batch_size, num_classes]
        return x

    # def forward(self, x):
    #     # x shape: (batch, 1, 28, 28)
    #     batch_size = x.shape[0]
    #     x = x.view(batch_size, -1, 7*7)
    #     x = self.patch_embedding(x)  # (batch, num_patches, dim)
    #     x = x + self.pos_embedding
    #     x = self.transformer(x)
    #     x = x.mean(dim=1)  # (batch, dim)
    #     x = self.mlp_head(x)
    #     return x