# PatchGAN discriminator for Phase 3.3 adversarial training: outputs a spatial patch logit map, not a scalar.
# Not in the model registry — instantiated directly inside train_vae when lambda_adversarial > 0.
import torch
import torch.nn as nn


def _init_weights(m):
    if isinstance(m, nn.Conv2d):
        nn.init.normal_(m.weight, 0.0, 0.02)
        if m.bias is not None:
            nn.init.zeros_(m.bias)


# Stride-2 + stride-1 conv chain → spatial patch logit map; supports image_size ∈ {64, 128}
# 64×64 → 6×6 map (70×70 receptive field); 128×128 adds an extra stride-2 layer. InstanceNorm except first layer.
class PatchGANDiscriminator(nn.Module):
    def __init__(self, ndf: int = 64, image_size: int = 64):
        super().__init__()
        layers: list[nn.Module] = [
            # First layer: no norm
            nn.Conv2d(3, ndf, 4, stride=2, padding=1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
        ]
        if image_size >= 128:
            layers += [
                nn.Conv2d(ndf, ndf, 4, stride=2, padding=1, bias=False),
                nn.InstanceNorm2d(ndf, affine=True),
                nn.LeakyReLU(0.2, inplace=True),
            ]
        layers += [
            nn.Conv2d(ndf, ndf * 2, 4, stride=2, padding=1, bias=False),
            nn.InstanceNorm2d(ndf * 2, affine=True),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(ndf * 2, ndf * 4, 4, stride=2, padding=1, bias=False),
            nn.InstanceNorm2d(ndf * 4, affine=True),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(ndf * 4, ndf * 8, 4, stride=1, padding=1, bias=False),
            nn.InstanceNorm2d(ndf * 8, affine=True),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(ndf * 8, 1, 4, stride=1, padding=1, bias=True),
        ]
        self.net = nn.Sequential(*layers)
        self.apply(_init_weights)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)  # (B, 1, H', W') — patch logit map


# Hinge loss for the discriminator (Lim & Ye, 2017)
def hinge_d_loss(real_logits: torch.Tensor, fake_logits: torch.Tensor) -> torch.Tensor:
    loss_real = torch.mean(torch.relu(1.0 - real_logits))
    loss_fake = torch.mean(torch.relu(1.0 + fake_logits))
    return 0.5 * (loss_real + loss_fake)


# Generator hinge loss — maximises D(fake)
def hinge_g_loss(fake_logits: torch.Tensor) -> torch.Tensor:
    return -torch.mean(fake_logits)
