import torch.nn.functional as F

class SelfAttention(nn.Module):
    def __init__(self, in_dim):
        super(SelfAttention, self).__init__()
        self.query_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim // 8, kernel_size=1)
        self.key_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim // 8, kernel_size=1)
        self.value_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim, kernel_size=1)
        self.gamma = nn.Parameter(torch.zeros(1))

    def forward(self, x):
        batch_size, C, width, height = x.size()
        proj_query = self.query_conv(x).view(batch_size, -1, width * height).permute(0, 2, 1)
        proj_key = self.key_conv(x).view(batch_size, -1, width * height)
        energy = torch.bmm(proj_query, proj_key)
        attention = F.softmax(energy, dim=-1)
        proj_value = self.value_conv(x).view(batch_size, -1, width * height)

        out = torch.bmm(proj_value, attention.permute(0, 2, 1))
        out = out.view(batch_size, C, width, height)

        out = self.gamma * out + x
        return out

class Decoder2(nn.Module):
    def __init__(self, embedding_dim=256):
        super(Decoder2, self).__init__()

        self.decoder = nn.Sequential(
            nn.Linear(768, 768),
            #nn.BatchNorm1d(768),
            nn.ReLU(),
            nn.Linear(768, 768),
            #nn.BatchNorm1d(768),
            nn.ReLU(),
            nn.Linear(768, 768),
            #nn.BatchNorm1d(768),
            nn.ReLU(),
            nn.Linear(768, 768),
            nn.ReLU(),
            nn.Linear(768, 768),nn.ReLU(),
            nn.Linear(768, 768),nn.ReLU(),
            nn.Linear(768, 768),nn.ReLU(),
            nn.Linear(768, 768),nn.ReLU(),
            nn.Linear(768, 768),nn.ReLU(),
            nn.Linear(768, 768),nn.ReLU(),
            nn.Linear(768, 768),
            #nn.BatchNorm1d(768),
            #nn.ReLU(),
            nn.Sigmoid()
        )

        hiddem_dim = 768
        latent_dim = 768

        self.fc2_mean = nn.Linear(hiddem_dim, latent_dim)
        self.fc2_logvar = nn.Linear(hiddem_dim, latent_dim)
        #self.decoder_input = nn.Linear(latent_dim, hiddem_dim)

    def encode(self, x):
        x = self.first_frame_conv(x)
        x = self.second_frame_conv(x)
        
        mu, logvar = torch.split(x, 512, dim=1)
        return mu, logvar
    
    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std
    
    def decode(self, z):
        #x = self.decoder_input(z)
        print(z.shape)
        x = self.decoder(z)
        return F.sigmoid(x)

    def forward(self, x):
        mu = self.fc2_mean(x)
        logvar = self.fc2_logvar(x)
        z = self.reparameterize(mu, logvar)
        z = self.decode(z)

        z = z.view(-1, 14, 14, 3, 16, 16)  # Reshape to [batch_size, 14, 14, 3, 16, 16]
        z = z.permute(0, 3, 1, 4, 2, 5)  # Permute to [batch_size, 3, 14, 16, 14, 16]
        z = z.contiguous().view(-1, 3, 224, 224)  # Reshape to [batch_size, 3, 224, 224]

        return z, mu, logvar
    
def loss_function(recon_x, x, mu, logvar):
    BCE = F.binary_cross_entropy(recon_x, x, reduction='sum')
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return BCE + KLD



почему он не делает цельную картинку а типо пиксель арт, типо он для каждого патча делает 1 цвет и не набор цветой, в 1 патче же 256 пикселей а он их все в 1 цвет делает