"""Trainable reverse of the generator, producing 64 sign logits for one label."""
import torch
from torch import nn
from torch.nn import functional as F


class ReverseResidualBlock(nn.Module):
    def __init__(self, channels_in: int, channels_out: int):
        super().__init__()
        self.conv2_t = nn.ConvTranspose2d(channels_in, channels_in, 3, padding=1)
        self.bn2 = nn.BatchNorm2d(channels_in)
        self.conv1_t = nn.ConvTranspose2d(channels_in, channels_out, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(channels_out)
        self.skip_t = nn.ConvTranspose2d(channels_in, channels_out, 1)

    def forward(self, x):
        h = self.bn2(F.leaky_relu(self.conv2_t(x), 0.1))
        h = F.avg_pool2d(self.conv1_t(h), 2)
        h = self.bn1(F.leaky_relu(h, 0.1))
        skip = F.avg_pool2d(self.skip_t(x), 2)
        return h + skip


class TransposedStegoAnalyzer(nn.Module):
    def __init__(self):
        super().__init__()
        self.output_conv_t = nn.ConvTranspose2d(1, 64, 3, padding=1)
        self.output_bn = nn.BatchNorm2d(64)
        self.reverse_up2 = ReverseResidualBlock(64, 128)
        self.reverse_up1 = ReverseResidualBlock(128, 256)
        self.fc_t = nn.Linear(256 * 7 * 7, 64)
        nn.init.normal_(self.fc_t.weight, std=0.005)
        nn.init.zeros_(self.fc_t.bias)

    def forward(self, image):
        if image.ndim != 4 or image.shape[1:] != (1, 28, 28):
            raise ValueError('Expected images of shape (batch, 1, 28, 28).')
        x = self.output_bn(F.leaky_relu(self.output_conv_t(image), 0.1))
        x = self.reverse_up2(x)
        x = self.reverse_up1(x)
        return self.fc_t(x.flatten(1))