"""Standalone inference architecture, copied unchanged from emnist_cgan64.py."""
import torch
from torch import nn
from torch.nn import functional as F

LATENT_DIM = 64


class ConditionalBatchNorm2d(nn.Module):
    """BN without affine + class-specific FiLM scaling/offset."""
    def __init__(self, channels: int, n_classes: int):
        super().__init__()
        self.bn = nn.BatchNorm2d(channels, affine=False)
        self.affine = nn.Embedding(n_classes, 2 * channels)
        nn.init.zeros_(self.affine.weight)

    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        gamma, beta = self.affine(y).chunk(2, dim=1)
        x = self.bn(x)
        return x * (1 + gamma[:, :, None, None]) + beta[:, :, None, None]


class ResBlockUp(nn.Module):
    def __init__(self, c_in: int, c_out: int, n_classes: int):
        super().__init__()
        self.bn1 = ConditionalBatchNorm2d(c_in, n_classes)
        self.bn2 = ConditionalBatchNorm2d(c_out, n_classes)
        self.conv1 = nn.Conv2d(c_in, c_out, 3, padding=1)
        self.conv2 = nn.Conv2d(c_out, c_out, 3, padding=1)
        self.skip = nn.Conv2d(c_in, c_out, 1)

    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        h = F.relu(self.bn1(x, y), inplace=False)
        h = F.interpolate(h, scale_factor=2, mode="nearest")
        h = self.conv1(h)
        h = self.conv2(F.relu(self.bn2(h, y), inplace=False))
        s = self.skip(F.interpolate(x, scale_factor=2, mode="nearest"))
        return h + s


class Generator(nn.Module):
    def __init__(self, n_classes: int):
        super().__init__()
        self.n_classes = n_classes
        self.fc = nn.Linear(LATENT_DIM, 256 * 7 * 7)
        self.up1 = ResBlockUp(256, 128, n_classes)
        self.up2 = ResBlockUp(128, 64, n_classes)
        self.bn = ConditionalBatchNorm2d(64, n_classes)
        self.out = nn.Conv2d(64, 1, kernel_size=3, padding=1)

    def forward(self, z: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        if z.ndim != 2 or z.shape[1] != LATENT_DIM:
            raise ValueError(f"z must be (batch,{LATENT_DIM}), got {tuple(z.shape)}")
        x = self.fc(z).reshape(-1, 256, 7, 7)
        x = self.up1(x, y)  # 128 x 14 x 14
        x = self.up2(x, y)  # 64 x 28 x 28
        return torch.tanh(self.out(F.relu(self.bn(x, y), inplace=False)))