"""LW-Vit"""

import torch 
import torch.nn as nn
import torch.nn.functional as F

class HSwish(nn.Module):
    """Hard Swish activation function"""
    def forward(self, x):
        out = x * F.relu6(x + 3, inplace=True) / 6
        return out

class MV2_Block(nn.Module):
    """MobileNetV2 Inverted Residual Block with h-swish activation"""
    def __init__(self, in_channels, expand_channels, out_channels, stride):
        super().__init__()
        self.stride = stride
        self.use_res_connect = self.stride == 1 and in_channels == out_channels
        
        self.expand = nn.Sequential(
            nn.Conv2d(in_channels, expand_channels, kernel_size=1, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(expand_channels),
            HSwish()
        )

        self.depthwise = nn.Sequential(
            nn.Conv2d(expand_channels, expand_channels, kernel_size=3, stride=stride, 
                     padding=1, groups=expand_channels, bias=False),
            nn.BatchNorm2d(expand_channels),
            HSwish()
        )

        self.project = nn.Sequential(
            nn.Conv2d(expand_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(out_channels)
        )

    def forward(self, x):
        if self.use_res_connect:
            return x + self.project(self.depthwise(self.expand(x)))
        else:
            return self.project(self.depthwise(self.expand(x)))
        
class LW_ViT_Transformer_Block(nn.Module):
    """Lightweight Vision Transformer Block"""
    def __init__(self, in_channels, out_channels, patch_size=2, heads=4, dim=128):
        super().__init__()
        self.patch_size = patch_size
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.dim = dim
        
        # Downsampling and channel adjustment
        self.downsample = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=1)
        self.channel_adjust = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1)
        
        # Calculate patch dimension
        patch_dim = out_channels * patch_size * patch_size
        # Add projection layer if needed
        self.projection = nn.Identity() if patch_dim == dim else nn.Linear(patch_dim, dim)
        
        # Transformer components
        self.norm1 = nn.LayerNorm(dim)
        self.attn = nn.MultiheadAttention(embed_dim=dim, num_heads=heads, batch_first=True)
        self.norm2 = nn.LayerNorm(dim)
        self.ffn = nn.Sequential(
            nn.Linear(dim, dim * 4),
            nn.GELU(),
            nn.Linear(dim * 4, dim)
        )
        
        # Final processing
        self.final_conv = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)
        self.norm_out = nn.BatchNorm2d(out_channels)
        self.act = HSwish()

    def forward(self, x):
        # Initial downsampling and channel adjustment
        x = self.downsample(x)
        x = self.channel_adjust(x)
        
        B, C, H, W = x.shape
        
        # Convert feature map to patches and embed
        # Reshape to [B, C, H/P, P, W/P, P]
        x_reshaped = x.reshape(B, C, H // self.patch_size, self.patch_size, 
                              W // self.patch_size, self.patch_size)
        # Permute to [B, H/P, W/P, C, P, P]
        x_permuted = x_reshaped.permute(0, 2, 4, 1, 3, 5)
        # Reshape to [B, H/P * W/P, C * P * P]
        patches = x_permuted.reshape(B, (H // self.patch_size) * (W // self.patch_size), -1)
        
        # Apply projection if needed
        patches = self.projection(patches)
        
        # Apply transformer operations
        normed_patches = self.norm1(patches)
        attn_out, _ = self.attn(normed_patches, normed_patches, normed_patches)
        patches = patches + attn_out
        
        normed_patches = self.norm2(patches)
        ffn_out = self.ffn(normed_patches)
        patches = patches + ffn_out
        
        # Reshape back to feature map
        # Reshape to [B, H/P, W/P, C, P, P]
        x_back = patches.reshape(B, H // self.patch_size, W // self.patch_size, 
                                C, self.patch_size, self.patch_size)
        # Permute to [B, C, H/P, P, W/P, P]
        x_back = x_back.permute(0, 3, 1, 4, 2, 5)
        # Reshape to [B, C, H, W]
        x_out = x_back.reshape(B, C, H, W)
        
        # Final processing
        out = self.final_conv(x_out)
        out = self.norm_out(out)
        out = self.act(out)
        
        return out

class LW_ViT(nn.Module):
    """Lightweight Vision Transformer for Chinese Character Recognition"""
    def __init__(self, base_channels=16, num_classes=3892):
        super().__init__()
        
        # Stem layer
        self.stem = nn.Sequential(
            nn.Conv2d(3, base_channels, 3, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(base_channels),
            HSwish()
        )

        # Feature extraction layers
        self.features = nn.Sequential(
            # MV2 blocks
            MV2_Block(base_channels, 4*base_channels, base_channels, stride=1),
            MV2_Block(base_channels, 4*base_channels, base_channels, stride=2),
            MV2_Block(base_channels, 4*base_channels, base_channels, stride=1),
            MV2_Block(base_channels, 4*base_channels, base_channels, stride=2),
            
            # Transformer block
            LW_ViT_Transformer_Block(
                in_channels=base_channels, 
                out_channels=2*base_channels, 
                patch_size=2, 
                heads=4, 
                dim=128
            ),
            
            # Final convolutional layer
            nn.Conv2d(2*base_channels, 4*base_channels, kernel_size=1, stride=1),
            nn.BatchNorm2d(4*base_channels),
            HSwish()
        )
        
        # Global pooling and classifier
        self.global_pool = nn.AdaptiveAvgPool2d(1)
        self.dropout = nn.Dropout(0.1)  # Added dropout for regularization
        
        self.classifier = nn.Sequential(
            nn.Linear(4*base_channels, 8*base_channels),
            nn.BatchNorm1d(8*base_channels),
            HSwish(),
            nn.Dropout(0.1),  # Added dropout for regularization
            nn.Linear(8*base_channels, num_classes)
        )

    def forward(self, x):
        x = self.stem(x)
        x = self.features(x)
        x = self.global_pool(x)
        x = torch.flatten(x, 1)
        x = self.dropout(x)
        x = self.classifier(x)
        return x
    

num_cls = 3892
num_epochs = 100
user_device = "cuda" if torch.cuda.is_available() else "cpu"
LW_model = LW_ViT(base_channels=16, num_classes=num_cls)