"""WGAN-GP variants.

wgan_basic  — Phase 2.2: BatchNorm/InstanceNorm, no attention, 64×64 only.
wgan        — Phase 2.3/2.4: GroupNorm/SpectralNorm + self-attention, size-agnostic.
"""
import math

import torch
import torch.nn as nn

from src.models import register


def _init_weights(m):
    if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
        nn.init.normal_(m.weight, 0.0, 0.02)
    elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)) and m.weight is not None:
        nn.init.normal_(m.weight, 1.0, 0.02)
        nn.init.zeros_(m.bias)


def _sn(module):
    return nn.utils.spectral_norm(module)


class SelfAttention(nn.Module):
    """SAGAN-style self-attention."""

    def __init__(self, in_ch: int):
        super().__init__()
        mid = max(in_ch // 8, 1)
        self.q = nn.Conv2d(in_ch, mid, 1, bias=False)
        self.k = nn.Conv2d(in_ch, mid, 1, bias=False)
        self.v = nn.Conv2d(in_ch, in_ch, 1, bias=False)
        self.gamma = nn.Parameter(torch.zeros(1))
        self._mid = mid

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b, c, h, w = x.shape
        q = self.q(x).view(b, self._mid, -1).transpose(-2, -1)  # (b, hw, mid)
        k = self.k(x).view(b, self._mid, -1)                    # (b, mid, hw)
        v = self.v(x).view(b, c, -1)                            # (b, c,  hw)
        attn = torch.softmax(q @ k * self._mid ** -0.5, dim=-1) # (b, hw, hw)
        return x + self.gamma * (v @ attn.transpose(-2, -1)).view(b, c, h, w)


# ---------------------------------------------------------------------------
# Phase 2.2 — basic WGAN-GP (BatchNorm in G, InstanceNorm in D, 64×64 only)
# ---------------------------------------------------------------------------

class WGANBasicGenerator(nn.Module):
    """Maps (latent_dim, 1, 1) -> (3, 64, 64) in [-1, 1].

    Same channel structure as DCGAN. BatchNorm in generator is fine because
    WGAN-GP's constraint targets the critic, not the generator.
    """

    def __init__(self, latent_dim: int = 128, ngf: int = 64):
        super().__init__()
        self.net = nn.Sequential(
            # 1×1 → 4×4
            nn.ConvTranspose2d(latent_dim, ngf * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(ngf * 8), nn.ReLU(True),
            # 4×4 → 8×8
            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 4), nn.ReLU(True),
            # 8×8 → 16×16
            nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 2), nn.ReLU(True),
            # 16×16 → 32×32
            nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf), nn.ReLU(True),
            # 32×32 → 64×64
            nn.ConvTranspose2d(ngf, 3, 4, 2, 1, bias=False),
            nn.Tanh(),
        )
        self.apply(_init_weights)

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


class WGANBasicCritic(nn.Module):
    """WGAN-GP critic (64×64). InstanceNorm instead of BatchNorm — BatchNorm
    breaks the per-sample Lipschitz constraint the gradient penalty enforces.
    """

    def __init__(self, ndf: int = 64):
        super().__init__()
        self.net = nn.Sequential(
            # 64×64 → 32×32 (no norm on first layer)
            nn.Conv2d(3, ndf, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, True),
            # 32×32 → 16×16
            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
            nn.InstanceNorm2d(ndf * 2, affine=True),
            nn.LeakyReLU(0.2, True),
            # 16×16 → 8×8
            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
            nn.InstanceNorm2d(ndf * 4, affine=True),
            nn.LeakyReLU(0.2, True),
            # 8×8 → 4×4
            nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False),
            nn.InstanceNorm2d(ndf * 8, affine=True),
            nn.LeakyReLU(0.2, True),
            # 4×4 → 1×1 (score, no sigmoid)
            nn.Conv2d(ndf * 8, 1, 4, 1, 0, bias=False),
        )
        self.apply(_init_weights)

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


# ---------------------------------------------------------------------------
# Phase 2.3 / 2.4 — advanced WGAN-GP (GroupNorm, SpectralNorm, attention)
# ---------------------------------------------------------------------------

class WGANGenerator(nn.Module):
    """GroupNorm generator with SAGAN self-attention.

    Supports image_size ∈ {64, 128}.
    Stem is always 1×1 → 4×4 → 8×8 → 16×16 (ngf×8 → ngf×4 → ngf×2 channels).
    Attention at 16×16 always; additional attention at 32×32 for 128×128.
    """

    def __init__(self, latent_dim: int = 128, ngf: int = 128, image_size: int = 64):
        super().__init__()
        if image_size not in (64, 128):
            raise ValueError(f"WGANGenerator supports image_size 64 or 128, got {image_size}")

        self._image_size = image_size

        self.stem = nn.Sequential(
            # 1×1 → 4×4
            nn.ConvTranspose2d(latent_dim, ngf * 8, 4, 1, 0, bias=False),
            nn.GroupNorm(8, ngf * 8), nn.ReLU(True),
            # 4×4 → 8×8
            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
            nn.GroupNorm(8, ngf * 4), nn.ReLU(True),
            # 8×8 → 16×16
            nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
            nn.GroupNorm(8, ngf * 2), nn.ReLU(True),
        )  # output: (ngf×2, 16, 16)
        self.attn16 = SelfAttention(ngf * 2)

        if image_size == 64:
            self.mid = None
            self.attn32 = None
            self.tail = nn.Sequential(
                # 16×16 → 32×32
                nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1, bias=False),
                nn.GroupNorm(8, ngf), nn.ReLU(True),
                # 32×32 → 64×64
                nn.ConvTranspose2d(ngf, 3, 4, 2, 1, bias=False),
                nn.Tanh(),
            )
        else:  # 128
            self.mid = nn.Sequential(
                # 16×16 → 32×32
                nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1, bias=False),
                nn.GroupNorm(8, ngf), nn.ReLU(True),
            )
            self.attn32 = SelfAttention(ngf)
            self.tail = nn.Sequential(
                # 32×32 → 64×64
                nn.ConvTranspose2d(ngf, ngf // 2, 4, 2, 1, bias=False),
                nn.GroupNorm(8, ngf // 2), nn.ReLU(True),
                # 64×64 → 128×128
                nn.ConvTranspose2d(ngf // 2, 3, 4, 2, 1, bias=False),
                nn.Tanh(),
            )
        self.apply(_init_weights)

    def forward(self, z: torch.Tensor) -> torch.Tensor:
        h = self.attn16(self.stem(z))
        if self.mid is not None:
            h = self.attn32(self.mid(h))
        return self.tail(h)


class WGANCritic(nn.Module):
    """SpectralNorm critic with SAGAN self-attention.

    Supports image_size ∈ {64, 128}.
    Attention at 16×16 always; additional attention at 32×32 for 128×128.
    """

    def __init__(self, ndf: int = 128, image_size: int = 64):
        super().__init__()
        if image_size not in (64, 128):
            raise ValueError(f"WGANCritic supports image_size 64 or 128, got {image_size}")

        self._image_size = image_size

        if image_size == 64:
            # Head: 64→32 (ndf//2)
            self.head = nn.Sequential(
                _sn(nn.Conv2d(3, ndf // 2, 4, 2, 1, bias=False)),
                nn.LeakyReLU(0.2, True),
            )
            self.attn32 = None
            # 32→16 (ndf)
            self.mid = nn.Sequential(
                _sn(nn.Conv2d(ndf // 2, ndf, 4, 2, 1, bias=False)),
                nn.LeakyReLU(0.2, True),
            )
            attn_ch = ndf
        else:  # 128
            # Head: 128→64 (ndf//4), 64→32 (ndf//2)
            self.head = nn.Sequential(
                _sn(nn.Conv2d(3, ndf // 4, 4, 2, 1, bias=False)),
                nn.LeakyReLU(0.2, True),
                _sn(nn.Conv2d(ndf // 4, ndf // 2, 4, 2, 1, bias=False)),
                nn.LeakyReLU(0.2, True),
            )
            self.attn32 = SelfAttention(ndf // 2)
            # 32→16 (ndf)
            self.mid = nn.Sequential(
                _sn(nn.Conv2d(ndf // 2, ndf, 4, 2, 1, bias=False)),
                nn.LeakyReLU(0.2, True),
            )
            attn_ch = ndf

        self.attn16 = SelfAttention(attn_ch)

        # Tail: 16×16 → 8×8 → 4×4 → score
        self.tail = nn.Sequential(
            _sn(nn.Conv2d(attn_ch, attn_ch * 2, 4, 2, 1, bias=False)),
            nn.LeakyReLU(0.2, True),
            _sn(nn.Conv2d(attn_ch * 2, attn_ch * 4, 4, 2, 1, bias=False)),
            nn.LeakyReLU(0.2, True),
            _sn(nn.Conv2d(attn_ch * 4, 1, 4, 1, 0, bias=False)),
        )
        self.apply(_init_weights)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = self.head(x)
        if self.attn32 is not None:
            h = self.attn32(h)
        h = self.attn16(self.mid(h))
        return self.tail(h).view(x.size(0))


def _build_basic(cfg: dict):
    return (
        WGANBasicGenerator(latent_dim=cfg.get("latent_dim", 128), ngf=cfg.get("ngf", 64)),
        WGANBasicCritic(ndf=cfg.get("ndf", 64)),
    )


def _build(cfg: dict):
    image_size = cfg.get("image_size", 64)
    return (
        WGANGenerator(latent_dim=cfg.get("latent_dim", 128), ngf=cfg.get("ngf", 128), image_size=image_size),
        WGANCritic(ndf=cfg.get("ndf", 128), image_size=image_size),
    )


register("wgan_basic", _build_basic, kind="wgan")
register("wgan", _build, kind="wgan")
