"""PatchGAN discriminator for Phase 3.3 (VQGAN-lite adversarial training).

Outputs a spatial patch map instead of a single scalar — each patch
predicts real/fake independently. Loss is the mean over all patches.

Not registered in the model registry; instantiated 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)


class PatchGANDiscriminator(nn.Module):
    """Stride-2 + stride-1 convolution chain → spatial patch logit map.

    Supports image_size ∈ {64, 128}. For 64×64 input the final map is 6×6
    (70×70 receptive field). For 128×128 an extra stride-2 layer is added.
    InstanceNorm everywhere except the first layer.
    """

    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


def hinge_d_loss(real_logits: torch.Tensor, fake_logits: torch.Tensor) -> torch.Tensor:
    """Hinge loss for the discriminator (Lim & Ye, 2017)."""
    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)


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