import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple, Optional
import math

# =====================================================
# 1. Patch Embedding (The Heart of ViT)
# =====================================================
class PatchEmbed(nn.Module):
    """
    Split image into patches → flatten → linear projection
    Input : (B, C, H, W)
    Output: (B, num_patches, embed_dim)
    """
    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
        super().__init__()
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = (img_size // patch_size) ** 2

        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        B, C, H, W = x.shape
        # (B, embed_dim, H/p, W/p) → (B, embed_dim, num_patches) → (B, num_patches, embed_dim)
        x = self.proj(x).flatten(2).transpose(1, 2)
        return x  # (B, N, D)


# =====================================================
# 2. Positional Embedding + Class Token
# =====================================================
class VisionTransformer(nn.Module):
    def __init__(
        self,
        img_size: int = 224,
        patch_size: int = 16,
        in_chans: int = 3,
        num_classes: int = 1000,
        embed_dim: int = 768,
        depth: int = 12,
        num_heads: int = 12,
        mlp_ratio: float = 4.0,
        qkv_bias: bool = True,
        drop_rate: float = 0.0,
        attn_drop_rate: float = 0.0,
        drop_path_rate: float = 0.1,
        norm_layer: nn.Module = nn.LayerNorm,
        use_abs_pos_emb: bool = True,
        use_cls_token: bool = True,
    ):
        super().__init__()
        self.num_classes = num_classes
        self.embed_dim = embed_dim
        self.use_cls_token = use_cls_token

        # Patch embedding
        self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim)
        num_patches = self.patch_embed.num_patches

        # Class token
        if use_cls_token:
            self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        else:
            self.cls_token = None

        # Positional embedding
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + (1 if use_cls_token else 0), embed_dim)) \
            if use_abs_pos_emb else None

        self.pos_drop = nn.Dropout(p=drop_rate)

        # Stochastic depth (drop path)
        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]

        # Transformer blocks
        self.blocks = nn.ModuleList([
            TransformerBlock(
                dim=embed_dim,
                num_heads=num_heads,
                mlp_ratio=mlp_ratio,
                qkv_bias=qkv_bias,
                drop=drop_rate,
                attn_drop=attn_drop_rate,
                drop_path=dpr[i],
                norm_layer=norm_layer,
            )
            for i in range(depth)
        ])

        self.norm = norm_layer(embed_dim)

        # Classifier head
        self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()

        # Weight init
        if self.pos_embed is not None:
            nn.init.trunc_normal_(self.pos_embed, std=0.02)
        if self.cls_token is not None:
            nn.init.trunc_normal_(self.cls_token, std=0.02)
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            nn.init.trunc_normal_(m.weight, std=0.02)
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def forward_features(self, x):
        B = x.shape[0]
        x = self.patch_embed(x)  # (B, N, D)

        # Add cls token
        if self.cls_token is not None:
            cls_tokens = self.cls_token.expand(B, -1, -1)
            x = torch.cat((cls_tokens, x), dim=1)

        # Add positional embedding
        if self.pos_embed is not None:
            x = x + self.pos_embed

        x = self.pos_drop(x)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x)
        return x

    def forward(self, x):
        x = self.forward_features(x)

        # Use cls token or mean pooling
        if self.cls_token is not None:
            x = x[:, 0]
        else:
            x = x.mean(dim=1)

        x = self.head(x)
        return x


# =====================================================
# 3. Core Transformer Block (Pre-LN + GELU + DropPath
# =====================================================
class TransformerBlock(nn.Module):
    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=True,
                 drop=0., attn_drop=0., drop_path=0., norm_layer=nn.LayerNorm):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias,
                              attn_drop=attn_drop, proj_drop=drop)
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim,
                       act_layer=nn.GELU, drop=drop)

    def forward(self, x):
        x = x + self.drop_path(self.attn(self.norm1(x)))
        x = x + self.drop_path(self.mlp(self.norm2(x)))
        return x


# =====================================================
# 4. Multi-Head Self Attention (Scaled Dot-Product)
# =====================================================
class Attention(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=True, attn_drop=0., proj_drop=0.):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = head_dim ** -0.5

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x):
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        q, k, v = qkv.unbind(0)  # (B, H, N, D)

        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x


# =====================================================
# 5. MLP + GELU + Dropout
# =====================================================
class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


# =====================================================
# 6. DropPath (Stochastic Depth)
# =====================================================
class DropPath(nn.Module):
    def __init__(self, drop_prob=None):
        super().__init__()
        self.drop_prob = drop_prob

    def forward(self, x):
        if self.drop_prob == 0. or not self.training:
            return x
        keep_prob = 1 - self.drop_prob
        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
        random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
        random_tensor.floor_()  # binarize
        output = x.div(keep_prob) * random_tensor
        return output


# =====================================================
# 7. Pre-built Models (Same as timm / HuggingFace)
# =====================================================
def vit_base_patch16_224(num_classes=1000):
    return VisionTransformer(
        img_size=224,
        patch_size=16,
        embed_dim=768,
        depth=12,
        num_heads=12,
        mlp_ratio=4,
        qkv_bias=True,
        num_classes=num_classes,
    )

def vit_large_patch16_224(num_classes=1000):
    return VisionTransformer(
        img_size=224,
        patch_size=16,
        embed_dim=1024,
        depth=24,
        num_heads=16,
        mlp_ratio=4,
        qkv_bias=True,
        num_classes=num_classes,
    )

def vit_huge_patch14_224(num_classes=1000):
    return VisionTransformer(
        img_size=224,
        patch_size=14,
        embed_dim=1280,
        depth=32,
        num_heads=16,
        mlp_ratio=4,
        qkv_bias=True,
        num_classes=num_classes,
    )


# =====================================================
# 8. Quick Test + CIFAR-10 / ImageNet Style
# =====================================================
if __name__ == "__main__":
    # Test on 224x224 ImageNet-like
    model = vit_base_patch16_224(num_classes=1000)
    x = torch.randn(2, 3, 224, 224)
    out = model(x)
    print(f"ViT-B/16 output: {out.shape}")  # → [2, 1000]

    # CIFAR-10 version (works perfectly)
    model_cifar = vit_base_patch16_224(num_classes=10)
    # Change patch embedding for 32x32
    model_cifar.patch_embed = PatchEmbed(img_size=32, patch_size=4, embed_dim=768)
    model_cifar.pos_embed = nn.Parameter(torch.zeros(1, (32//4)**2 + 1, 768))
    x_cifar = torch.randn(8, 3, 32, 32)
    print("CIFAR-10 ViT output:", model_cifar(x_cifar).shape)  # → [8, 10]

# 1. Use Relative Position Bias (Swin Transformer style)
# 2. Use Rotary Embeddings (RoPE) – used in Llama 3, Grok
# 3. Use LayerScale (CaiT)
# 4. Use GELU → SwiGLU (better performance)
# 5. Add Class-Attention (CAiT) or Token Labeling

import torchvision, torchvision.transforms as T
from torch.optim import AdamW

model = vit_base_patch16_224(num_classes=10)
model.patch_embed = PatchEmbed(32, 4, 3, 384)  # smaller model for CIFAR
model.pos_embed = nn.Parameter(torch.zeros(1, 65, 384))

transform = T.Compose([T.Resize(32), T.ToTensor(), T.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
loader = torch.utils.data.DataLoader(trainset, batch_size=256, shuffle=True)

optimizer = AdamW(model.parameters(), lr=3e-4, weight_decay=0.05)
model.cuda()

for epoch in range(5):
    for x, y in loader:
        x, y = x.cuda(), y.cuda()
        loss = F.cross_entropy(model(x), y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    print(f"Epoch {epoch} loss: {loss.item():.4f}")