import torch.nn as nn
from transformers import ViTModel

class SEMViTAutoencoder(nn.Module):
    def __init__(self):
        super().__init__()
        
        # Load pre-trained ViT-Base as encoder
        self.encoder = ViTModel.from_pretrained(
            'google/vit-base-patch16-224',  # ViT expects 224p, 3 channels (RGB)
            add_pooling_layer=False 
        )
        
               
        # Custom convolutional decoder to reconstruct the image
        # Decoder: Mapping ViT features (768) back to image space
        # 512/16 = 32 patches. Latent grid is 32x32.
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(768, 256, kernel_size=4, stride=2, padding=1), # 64x64
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1), # 128x128
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1), # 256x256
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.ConvTranspose2d(64, 3, kernel_size=4, stride=2, padding=1),   # 512x512
            nn.Sigmoid() # Output pixels in [0, 1]
        )

    def forward(self, x):
        # ViT requires position interpolation for resolutions != 224
        outputs = self.encoder(x, interpolate_pos_encoding=True)
        # Sequence to Grid: [Batch, 1025, 768] -> [Batch, 768, 32, 32]
        latent = outputs.last_hidden_state[:, 1:, :].transpose(1, 2).reshape(-1, 768, 32, 32)
        
        return self.decoder(latent)