import torch
import torch.nn as nn

class Encoder(nn.Module):
    def __init__(self, embedding_dim=256):
        super(Encoder, self).__init__()
        # Input channels: 6 (concatenated frames)
        self.features = nn.Sequential(
            nn.Conv2d(6, 64, kernel_size=4, stride=2, padding=1),   # Output: 64 x 256 x 256
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),  # Output: 128 x 128 x 128
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1),  # Output: 256 x 64 x 64
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1),  # Output: 512 x 32 x 32
            nn.BatchNorm2d(512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(512, 512, kernel_size=4, stride=2, padding=1),  # Output: 512 x 16 x 16
            nn.BatchNorm2d(512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.AdaptiveAvgPool2d((1, 1)),  # Output: 512 x 1 x 1
        )
        # Linear layer to produce the embedding
        self.fc = nn.Linear(512, embedding_dim)

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)  # Flatten to [batch_size, 512]
        embedding = self.fc(x)     # Output: [batch_size, embedding_dim]
        return embedding

class Decoder(nn.Module):
    def __init__(self, embedding_dim=256):
        super(Decoder, self).__init__()
        # Project embedding and expand to feature map
        self.fc = nn.Sequential(
            nn.Linear(embedding_dim, 512 * 4 * 4),
            nn.ReLU(True)
        )
        # Initial convolution to process the first frame
        self.first_frame_conv = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=7, stride=1, padding=3),  # Output: 64 x 512 x 512
            nn.BatchNorm2d(64),
            nn.ReLU(True),
            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),  # Output: 128 x 256 x 256
            nn.BatchNorm2d(128),
            nn.ReLU(True),
            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1),  # Output: 256 x 128 x 128
            nn.BatchNorm2d(256),
            nn.ReLU(True),
            nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1),  # Output: 512 x 64 x 64
            nn.BatchNorm2d(512),
            nn.ReLU(True),
            nn.Conv2d(512, 512, kernel_size=4, stride=2, padding=1),  # Output: 512 x 32 x 32
            nn.BatchNorm2d(512),
            nn.ReLU(True),
            nn.Conv2d(512, 512, kernel_size=4, stride=2, padding=1),  # Output: 512 x 16 x 16
            nn.BatchNorm2d(512),
            nn.ReLU(True),
            nn.AdaptiveAvgPool2d((4, 4)),  # Output: 512 x 4 x 4
        )
        # Decoder network
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(1024, 512, kernel_size=4, stride=2, padding=1),  # Output: 512 x 8 x 8
            nn.BatchNorm2d(512),
            nn.ReLU(True),
            nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1),  # Output: 256 x 16 x 16
            nn.BatchNorm2d(256),
            nn.ReLU(True),
            nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1),  # Output: 128 x 32 x 32
            nn.BatchNorm2d(128),
            nn.ReLU(True),
            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),  # Output: 64 x 64 x 64
            nn.BatchNorm2d(64),
            nn.ReLU(True),
            nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1),  # Output: 32 x 128 x 128
            nn.BatchNorm2d(32),
            nn.ReLU(True),
            nn.ConvTranspose2d(32, 16, kernel_size=4, stride=2, padding=1),  # Output: 16 x 256 x 256
            nn.BatchNorm2d(16),
            nn.ReLU(True),
            nn.ConvTranspose2d(16, 3, kernel_size=4, stride=2, padding=1),  # Output: 3 x 512 x 512
            nn.Tanh()  # Output values between -1 and 1
        ) # torch.Size([32, 3, 128, 128])

    def forward(self, first_frame, embedding):
        # Process first frame
        first_frame_features = self.first_frame_conv(first_frame)  # Output: [batch_size, 512, 4, 4]
        # Process embedding
        embedding_features = self.fc(embedding)
        embedding_features = embedding_features.view(-1, 512, 4, 4)
        # Concatenate along channel dimension
        x = torch.cat((first_frame_features, embedding_features), dim=1)  # Output: [batch_size, 1024, 4, 4]
        x = self.decoder(x)
        #print(x.shape)
        return x
    

    


подумай об этой архитектуре и улучши ее, сейчас она даже не может адекватно реконструировать 2 фрейм, хотя они почти одинаковые, проблема не в датасете если что, он оч большой