"""Vanilla DCGAN (Radford et al., 2015).

Used as the Phase 1 baseline for cheap pipeline ablations. Architecture is
intentionally minimal — BatchNorm in both networks, no spectral norm, no
attention, no gradient penalty. The whole point is to be the cheapest GAN
we can run, so 1A–1D pipeline deltas show up in FID quickly.

Depth scales with image_size: each step doubles the spatial dimension,
starting from 4×4 after the first transposed conv.
    64  -> 5 layers (1 -> 4 -> 8 -> 16 -> 32 -> 64)
    128 -> 6 layers (1 -> 4 -> 8 -> 16 -> 32 -> 64 -> 128)
"""
import math

import torch
import torch.nn as nn

from src.models import register


def _init_weights(m):
    classname = m.__class__.__name__
    if "Conv" in classname:
        nn.init.normal_(m.weight, 0.0, 0.02)
    elif "BatchNorm" in classname:
        nn.init.normal_(m.weight, 1.0, 0.02)
        nn.init.zeros_(m.bias)


def _n_upsamples(image_size: int) -> int:
    """Number of 2x upsampling steps from 4x4 to image_size."""
    if image_size < 8 or image_size & (image_size - 1):
        raise ValueError(f"image_size must be a power of two ≥ 8, got {image_size}")
    return int(math.log2(image_size)) - 2  # 64 -> 4, 128 -> 5


class DCGANGenerator(nn.Module):
    """Maps (latent_dim x 1 x 1) -> (3 x image_size x image_size) in [-1, 1]."""

    def __init__(self, latent_dim: int = 100, ngf: int = 64, image_size: int = 64):
        super().__init__()
        n_up = _n_upsamples(image_size)            # 64 -> 4 upsamples after the 1->4 init
        max_mult = 2 ** (n_up - 1)                 # channel multiplier at the 4x4 stage

        layers: list[nn.Module] = [
            # 1x1 -> 4x4
            nn.ConvTranspose2d(latent_dim, ngf * max_mult, 4, 1, 0, bias=False),
            nn.BatchNorm2d(ngf * max_mult),
            nn.ReLU(inplace=True),
        ]
        # Each step halves the channel multiplier and doubles spatial size.
        mult = max_mult
        for _ in range(n_up - 1):
            layers += [
                nn.ConvTranspose2d(ngf * mult, ngf * mult // 2, 4, 2, 1, bias=False),
                nn.BatchNorm2d(ngf * mult // 2),
                nn.ReLU(inplace=True),
            ]
            mult //= 2
        # Final layer to 3 channels, no BN, Tanh.
        layers += [
            nn.ConvTranspose2d(ngf * mult, 3, 4, 2, 1, bias=False),
            nn.Tanh(),
        ]
        self.net = nn.Sequential(*layers)
        self.apply(_init_weights)

    def forward(self, z: torch.Tensor) -> torch.Tensor:
        return self.net(z)


class DCGANDiscriminator(nn.Module):
    """Maps (3 x image_size x image_size) -> scalar logit (no sigmoid)."""

    def __init__(self, ndf: int = 64, image_size: int = 64):
        super().__init__()
        n_down = _n_upsamples(image_size)
        layers: list[nn.Module] = [
            # First layer: no BN
            nn.Conv2d(3, ndf, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
        ]
        mult = 1
        for _ in range(n_down - 1):
            layers += [
                nn.Conv2d(ndf * mult, ndf * mult * 2, 4, 2, 1, bias=False),
                nn.BatchNorm2d(ndf * mult * 2),
                nn.LeakyReLU(0.2, inplace=True),
            ]
            mult *= 2
        # 4x4 -> 1x1, scalar logit
        layers += [nn.Conv2d(ndf * mult, 1, 4, 1, 0, bias=False)]
        self.net = nn.Sequential(*layers)
        self.apply(_init_weights)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x).view(x.size(0))


def _build(cfg: dict):
    image_size = cfg.get("image_size", 64)
    return (
        DCGANGenerator(
            latent_dim=cfg.get("latent_dim", 100),
            ngf=cfg.get("ngf", 64),
            image_size=image_size,
        ),
        DCGANDiscriminator(
            ndf=cfg.get("ndf", 64),
            image_size=image_size,
        ),
    )


register("dcgan", _build, kind="dcgan")
