"""Convolutional VAE for Phase 3.

Encoder uses stride-2 Conv → flatten → linear (μ, log σ²).
Decoder uses Linear → Upsample(nearest) + Conv to avoid ConvTranspose2d
checkerboard artefacts.

Registered as kind="vae". The run.py dispatcher passes the model to
train_vae(), which internally builds perceptual loss and PatchGAN when
the corresponding lambdas are non-zero.
"""
import math

import torch
import torch.nn as nn
import torch.nn.functional as F

from src.models import register


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


def _upsample_block(in_ch: int, out_ch: int) -> nn.Sequential:
    """Nearest-neighbour upsample followed by a 3×3 conv — no checkerboard."""
    return nn.Sequential(
        nn.Upsample(scale_factor=2, mode="nearest"),
        nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False),
        nn.BatchNorm2d(out_ch),
        nn.ReLU(inplace=True),
    )


class VAE(nn.Module):
    """Convolutional VAE. image_size must be a power-of-two ≥ 32.

    Spatial bottleneck is always at 4×4 regardless of image_size —
    the encoder and decoder scale the number of stride-2 steps accordingly.
    """

    def __init__(self, latent_dim: int = 256, ngf: int = 64, image_size: int = 64):
        super().__init__()
        if image_size < 32 or (image_size & (image_size - 1)):
            raise ValueError(f"image_size must be a power-of-two ≥ 32, got {image_size}")

        self.latent_dim = latent_dim
        self.image_size = image_size

        n_down = int(math.log2(image_size)) - 2  # steps from image_size to 4×4
        # 64 → n_down=4: 64→32→16→8→4
        # 128 → n_down=5: 128→64→32→16→8→4

        # ── Encoder ──────────────────────────────────────────────────────────
        enc_layers: list[nn.Module] = [
            nn.Conv2d(3, ngf, 4, stride=2, padding=1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
        ]
        ch = ngf
        for _ in range(n_down - 1):
            enc_layers += [
                nn.Conv2d(ch, ch * 2, 4, stride=2, padding=1, bias=False),
                nn.BatchNorm2d(ch * 2),
                nn.LeakyReLU(0.2, inplace=True),
            ]
            ch *= 2
        # ch = ngf * 2^(n_down-1); spatial = 4×4
        self.encoder = nn.Sequential(*enc_layers)

        flat = ch * 4 * 4
        self.fc_mu  = nn.Linear(flat, latent_dim)
        self.fc_lv  = nn.Linear(flat, latent_dim)

        # ── Decoder ──────────────────────────────────────────────────────────
        self.fc_dec = nn.Linear(latent_dim, flat)
        self._dec_ch = ch  # channels at the 4×4 bottleneck

        dec_layers: list[nn.Module] = []
        for _ in range(n_down - 1):
            dec_layers.append(_upsample_block(ch, ch // 2))
            ch //= 2
        # Final upsample to image_size, output 3 channels, no BN, Tanh
        dec_layers += [
            nn.Upsample(scale_factor=2, mode="nearest"),
            nn.Conv2d(ch, 3, 3, padding=1, bias=True),
            nn.Tanh(),
        ]
        self.decoder = nn.Sequential(*dec_layers)

        self.apply(_init_weights)

    # ── Interface ────────────────────────────────────────────────────────────

    def encode(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        h = self.encoder(x).flatten(1)
        log_var = self.fc_lv(h).clamp(-10.0, 10.0)
        return self.fc_mu(h), log_var

    def reparameterize(self, mu: torch.Tensor, log_var: torch.Tensor) -> torch.Tensor:
        std = torch.exp(0.5 * log_var)
        return mu + std * torch.randn_like(std)

    def decode(self, z: torch.Tensor) -> torch.Tensor:
        h = self.fc_dec(z).view(z.size(0), self._dec_ch, 4, 4)
        return self.decoder(h)

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Returns (reconstruction, mu, log_var)."""
        mu, log_var = self.encode(x)
        z = self.reparameterize(mu, log_var)
        return self.decode(z), mu, log_var

    @torch.no_grad()
    def sample(self, n: int, device) -> torch.Tensor:
        """Sample n images by drawing z ~ N(0, I)."""
        z = torch.randn(n, self.latent_dim, device=device)
        return self.decode(z)


def _build(cfg: dict):
    return VAE(
        latent_dim=cfg.get("latent_dim", 256),
        ngf=cfg.get("ngf", 64),
        image_size=cfg.get("image_size", 64),
    )


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