"""
Upgraded UNet with self-attention + VQGAN latent-space wrapper.

Two modes:
  1. Pixel-space (no VQGAN): use BBDMUNet directly, same as before but with attention
  2. Latent-space (with VQGAN): use LatentBBDM which wraps encode/decode around BBDMUNet

Usage:
    # Pixel-space (drop-in replacement for SimpleUNet)
    model = SimpleUNet(in_channels=3, base_channels=64, image_size=256)
    noise_pred = model(x_t, x_T, t)

    # Latent-space (requires pretrained VQGAN)
    latent_model = LatentBBDM(
        unet_channels=64,
        latent_channels=4,     # VQGAN latent dim (usually 3 or 4)
        vqgan_ckpt="path/to/vqgan.ckpt",
        vqgan_config="path/to/vqgan_config.yaml",
    )
"""

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


# ============================================================
# Building blocks
# ============================================================

class SinusoidalPosEmb(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, t):
        half = self.dim // 2
        emb = math.log(10000) / (half - 1)
        emb = torch.exp(torch.arange(half, device=t.device) * -emb)
        emb = t.float().unsqueeze(1) * emb.unsqueeze(0)
        return torch.cat([emb.sin(), emb.cos()], dim=-1)


class ResBlock(nn.Module):
    """Residual block with timestep conditioning via addition."""

    def __init__(self, in_ch, out_ch, t_dim, dropout=0.0):
        super().__init__()
        self.norm1 = nn.GroupNorm(min(32, in_ch), in_ch)
        self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)
        self.t_proj = nn.Linear(t_dim, out_ch)
        self.norm2 = nn.GroupNorm(min(32, out_ch), out_ch)
        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)
        self.skip = nn.Conv2d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity()

    def forward(self, x, t_emb):
        h = self.conv1(F.silu(self.norm1(x)))
        h = h + self.t_proj(F.silu(t_emb)).unsqueeze(-1).unsqueeze(-1)
        h = self.conv2(self.dropout(F.silu(self.norm2(h))))
        return h + self.skip(x)


class SelfAttention(nn.Module):
    """
    Multi-head self-attention for spatial feature maps.
    Applied at 16x16 and 32x32 resolutions to capture global context.
    """

    def __init__(self, channels, num_heads=4):
        super().__init__()
        self.channels = channels
        self.num_heads = num_heads
        assert channels % num_heads == 0

        self.norm = nn.GroupNorm(min(32, channels), channels)
        self.qkv = nn.Conv1d(channels, channels * 3, 1)
        self.proj = nn.Conv1d(channels, channels, 1)
        self.scale = (channels // num_heads) ** -0.5

    def forward(self, x):
        B, C, H, W = x.shape
        h = self.norm(x).view(B, C, H * W)

        qkv = self.qkv(h).view(B, 3, self.num_heads, C // self.num_heads, H * W)
        q, k, v = qkv[:, 0], qkv[:, 1], qkv[:, 2]

        q = q.permute(0, 1, 3, 2)  # [B, heads, HW, dim]
        k = k.permute(0, 1, 2, 3)  # [B, heads, dim, HW]
        v = v.permute(0, 1, 3, 2)  # [B, heads, HW, dim]

        attn = torch.matmul(q, k) * self.scale
        attn = attn.softmax(dim=-1)
        out = torch.matmul(attn, v)
        out = out.permute(0, 1, 3, 2).contiguous().view(B, C, H * W)
        out = self.proj(out).view(B, C, H, W)

        return x + out


class CrossAttention(nn.Module):
    """
    Multi-head cross-attention for spatial feature maps.
    Queries come from x_t trunk, keys/values come from y trunk.
    """

    def __init__(self, channels, num_heads=4):
        super().__init__()
        self.channels = channels
        self.num_heads = num_heads
        assert channels % num_heads == 0

        self.norm_q = nn.GroupNorm(min(32, channels), channels)
        self.norm_kv = nn.GroupNorm(min(32, channels), channels)
        self.to_q = nn.Conv1d(channels, channels, 1)
        self.to_kv = nn.Conv1d(channels, channels * 2, 1)
        self.proj = nn.Conv1d(channels, channels, 1)
        self.scale = (channels // num_heads) ** -0.5

    def forward(self, x, context):
        """
        Args:
            x: query features [B, C, H, W] from x_t trunk
            context: key/value features [B, C, H, W] from y trunk
        """
        B, C, H, W = x.shape
        q_in = self.norm_q(x).view(B, C, H * W)
        kv_in = self.norm_kv(context).view(B, C, H * W)

        q = self.to_q(q_in).view(B, self.num_heads, C // self.num_heads, H * W)
        kv = self.to_kv(kv_in).view(B, 2, self.num_heads, C // self.num_heads, H * W)
        k, v = kv[:, 0], kv[:, 1]

        q = q.permute(0, 1, 3, 2)   # [B, heads, HW, dim]
        k = k.permute(0, 1, 2, 3)   # [B, heads, dim, HW]
        v = v.permute(0, 1, 3, 2)   # [B, heads, HW, dim]

        attn = torch.matmul(q, k) * self.scale
        attn = attn.softmax(dim=-1)
        out = torch.matmul(attn, v)
        out = out.permute(0, 1, 3, 2).contiguous().view(B, C, H * W)
        out = self.proj(out).view(B, C, H, W)

        return x + out


class Downsample(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv = nn.Conv2d(channels, channels, 3, stride=2, padding=1)

    def forward(self, x):
        return self.conv(x)


class Upsample(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv = nn.Conv2d(channels, channels, 3, padding=1)

    def forward(self, x):
        x = F.interpolate(x, scale_factor=2, mode="nearest")
        return self.conv(x)


# ============================================================
# Main UNet with attention
# ============================================================

class BBDMUNet(nn.Module):
    """
    UNet for BBDM with self-attention at specified resolutions.

    Changes from SimpleUNet:
      - Self-attention at 16x16 and 32x32
      - Norm-before-conv ordering (DDPM paper convention)
      - Interpolation upsampling (fewer checkerboard artifacts)
      - Zero-initialized output conv
      - Larger timestep embedding (256-dim with hidden expansion)
    """

    def __init__(
        self,
        in_channels=3,
        base_channels=64,
        t_dim=256,
        image_size=256,
        channel_mults=(1, 2, 4, 4),
        attn_resolutions=(16, 32),
        num_heads=4,
        dropout=0.0,
        dual_output=False,
        use_x1_cond=True,
    ):
        super().__init__()
        self.image_size = image_size
        self.in_channels = in_channels
        self.dual_output = dual_output
        self.use_x1_cond = use_x1_cond

        # Timestep embedding
        self.time_mlp = nn.Sequential(
            SinusoidalPosEmb(t_dim),
            nn.Linear(t_dim, t_dim * 4),
            nn.SiLU(),
            nn.Linear(t_dim * 4, t_dim),
        )

        ch = base_channels

        # Input projection: 6-ch (x_t + x_1) when use_x1_cond=True, else 3-ch (x_t alone)
        n_in = in_channels * (2 if use_x1_cond else 1)
        self.input_conv = nn.Conv2d(n_in, ch, 3, padding=1)

        # ---- Encoder ----
        self.enc_blocks = nn.ModuleList()
        self.enc_attns = nn.ModuleList()
        self.enc_downs = nn.ModuleList()

        skip_channels = [ch]
        cur_ch = ch
        cur_res = image_size

        for i, mult in enumerate(channel_mults):
            out_ch = ch * mult
            self.enc_blocks.append(ResBlock(cur_ch, out_ch, t_dim, dropout))
            cur_ch = out_ch
            skip_channels.append(cur_ch)

            if cur_res in attn_resolutions:
                self.enc_attns.append(SelfAttention(cur_ch, num_heads))
            else:
                self.enc_attns.append(nn.Identity())

            if i < len(channel_mults) - 1:
                self.enc_downs.append(Downsample(cur_ch))
                cur_res //= 2
            else:
                self.enc_downs.append(nn.Identity())

        # ---- Bottleneck ----
        self.mid1 = ResBlock(cur_ch, cur_ch, t_dim, dropout)
        self.mid_attn = SelfAttention(cur_ch, num_heads)
        self.mid2 = ResBlock(cur_ch, cur_ch, t_dim, dropout)

        # ---- Decoder ----
        self.dec_blocks = nn.ModuleList()
        self.dec_attns = nn.ModuleList()
        self.dec_ups = nn.ModuleList()

        for i in reversed(range(len(channel_mults))):
            mult = channel_mults[i]
            out_ch = ch * mult
            skip_ch = skip_channels.pop()

            self.dec_blocks.append(ResBlock(cur_ch + skip_ch, out_ch, t_dim, dropout))
            cur_ch = out_ch

            dec_res = image_size // (2 ** i) if i < len(channel_mults) - 1 else cur_res
            if dec_res in attn_resolutions:
                self.dec_attns.append(SelfAttention(cur_ch, num_heads))
            else:
                self.dec_attns.append(nn.Identity())

            if i > 0:
                self.dec_ups.append(Upsample(cur_ch))
            else:
                self.dec_ups.append(nn.Identity())

        # ---- Output ----
        self.out_norm = nn.GroupNorm(min(32, cur_ch), cur_ch)
        if dual_output:
            # Two parallel heads: predicts noise z and clean image x_0
            self.out_z = nn.Conv2d(cur_ch, in_channels, 3, padding=1)
            self.out_x0 = nn.Conv2d(cur_ch, in_channels, 3, padding=1)
            nn.init.zeros_(self.out_z.weight)
            nn.init.zeros_(self.out_z.bias)
            nn.init.zeros_(self.out_x0.weight)
            nn.init.zeros_(self.out_x0.bias)
        else:
            self.out_conv = nn.Conv2d(cur_ch, in_channels, 3, padding=1)
            nn.init.zeros_(self.out_conv.weight)
            nn.init.zeros_(self.out_conv.bias)

    def forward(self, x_t, x_T, t, cond_mask=None):
        """
        Args:
            x_t: noisy intermediate [B, C, H, W]
            x_T: source conditioning [B, C, H, W]
            t: timestep [B]
            cond_mask: optional [B, 1, 1, 1] binary mask. 0 = drop conditioning (for CFG training).

        Returns:
            single tensor [B, C, H, W] when dual_output=False (default),
            tuple (z_pred, x0_pred) each [B, C, H, W] when dual_output=True.
        """
        t_emb = self.time_mlp(t)
        if self.use_x1_cond:
            if cond_mask is not None:
                x_T = x_T * cond_mask
            x = torch.cat([x_t, x_T], dim=1)
        else:
            x = x_t   # x_1 NOT given as explicit input; only enters via the bridge state x_t
        h = self.input_conv(x)

        # Encoder with skip connections
        skips = [h]
        for block, attn, down in zip(self.enc_blocks, self.enc_attns, self.enc_downs):
            h = block(h, t_emb)
            h = attn(h)
            skips.append(h)
            h = down(h)

        # Bottleneck
        h = self.mid1(h, t_emb)
        h = self.mid_attn(h)
        h = self.mid2(h, t_emb)

        # Decoder
        for block, attn, up in zip(self.dec_blocks, self.dec_attns, self.dec_ups):
            h = torch.cat([h, skips.pop()], dim=1)
            h = block(h, t_emb)
            h = attn(h)
            h = up(h)

        h = F.silu(self.out_norm(h))
        if self.dual_output:
            return self.out_z(h), self.out_x0(h)
        return self.out_conv(h)


# ============================================================
# Dual-Trunk UNet (cross-attention conditioning)
# ============================================================

class DualTrunkBBDMUNet(nn.Module):
    """
    Dual-trunk UNet for BBDM with cross-attention conditioning.

    y-trunk (encoder only): Extracts multi-resolution features from source image y.
    x_t-trunk (full UNet): Denoises x_t, cross-attending to y-trunk features at each
    resolution level that has attention.

    CFG: During training, y-trunk features are randomly replaced with zeros.
    At inference, the model handles CFG internally via encode_y() + forward().

    Inspired by TryOnDiffusion's parallel UNet architecture.
    """

    def __init__(
        self,
        in_channels=3,
        base_channels=64,
        t_dim=256,
        image_size=256,
        channel_mults=(1, 2, 4, 4),
        attn_resolutions=(16, 32),
        num_heads=4,
        dropout=0.0,
    ):
        super().__init__()
        self.image_size = image_size
        self.in_channels = in_channels
        self.is_dual_trunk = True  # flag for bridge/training code to detect

        ch = base_channels

        # ---- Timestep embedding (only for x_t trunk) ----
        self.time_mlp = nn.Sequential(
            SinusoidalPosEmb(t_dim),
            nn.Linear(t_dim, t_dim * 4),
            nn.SiLU(),
            nn.Linear(t_dim * 4, t_dim),
        )

        # ================================================================
        # Y-TRUNK (encoder only, no timestep, extracts conditioning features)
        # ================================================================
        self.y_input_conv = nn.Conv2d(in_channels, ch, 3, padding=1)

        self.y_enc_blocks = nn.ModuleList()
        self.y_enc_attns = nn.ModuleList()
        self.y_enc_downs = nn.ModuleList()

        cur_ch = ch
        cur_res = image_size
        # Track y-trunk output channels at each level for cross-attention
        self._y_channels = [cur_ch]  # after input conv

        for i, mult in enumerate(channel_mults):
            out_ch = ch * mult
            # y-trunk uses ResBlock but with a dummy t_dim — we'll pass zeros
            self.y_enc_blocks.append(ResBlock(cur_ch, out_ch, t_dim, dropout))
            cur_ch = out_ch
            self._y_channels.append(cur_ch)

            if cur_res in attn_resolutions:
                self.y_enc_attns.append(SelfAttention(cur_ch, num_heads))
            else:
                self.y_enc_attns.append(nn.Identity())

            if i < len(channel_mults) - 1:
                self.y_enc_downs.append(Downsample(cur_ch))
                cur_res //= 2
            else:
                self.y_enc_downs.append(nn.Identity())

        # Y-trunk bottleneck
        self.y_mid1 = ResBlock(cur_ch, cur_ch, t_dim, dropout)
        self.y_mid_attn = SelfAttention(cur_ch, num_heads)
        self.y_mid2 = ResBlock(cur_ch, cur_ch, t_dim, dropout)

        # ================================================================
        # X_T TRUNK (full UNet with cross-attention to y-trunk features)
        # ================================================================
        cur_ch = ch
        cur_res = image_size

        self.input_conv = nn.Conv2d(in_channels, ch, 3, padding=1)  # 3ch, NOT 6ch

        # ---- Encoder ----
        self.enc_blocks = nn.ModuleList()
        self.enc_self_attns = nn.ModuleList()
        self.enc_cross_attns = nn.ModuleList()
        self.enc_downs = nn.ModuleList()

        skip_channels = [ch]

        for i, mult in enumerate(channel_mults):
            out_ch = ch * mult
            self.enc_blocks.append(ResBlock(cur_ch, out_ch, t_dim, dropout))
            cur_ch = out_ch
            skip_channels.append(cur_ch)

            if cur_res in attn_resolutions:
                self.enc_self_attns.append(SelfAttention(cur_ch, num_heads))
                self.enc_cross_attns.append(CrossAttention(cur_ch, num_heads))
            else:
                self.enc_self_attns.append(nn.Identity())
                self.enc_cross_attns.append(None)  # placeholder

            if i < len(channel_mults) - 1:
                self.enc_downs.append(Downsample(cur_ch))
                cur_res //= 2
            else:
                self.enc_downs.append(nn.Identity())

        # Use ModuleList for proper registration of cross-attn (replace None with Identity)
        # We handle None in forward manually
        self._enc_cross_attn_indices = []
        for i, ca in enumerate(self.enc_cross_attns):
            if ca is not None:
                self._enc_cross_attn_indices.append(i)
        # Re-register as ModuleList (only non-None)
        self.enc_cross_attns = nn.ModuleList(
            [ca if ca is not None else nn.Identity() for ca in self.enc_cross_attns]
        )

        # ---- Bottleneck ----
        self.mid1 = ResBlock(cur_ch, cur_ch, t_dim, dropout)
        self.mid_self_attn = SelfAttention(cur_ch, num_heads)
        self.mid_cross_attn = CrossAttention(cur_ch, num_heads)
        self.mid2 = ResBlock(cur_ch, cur_ch, t_dim, dropout)

        # ---- Decoder ----
        self.dec_blocks = nn.ModuleList()
        self.dec_self_attns = nn.ModuleList()
        self.dec_cross_attns = nn.ModuleList()
        self.dec_ups = nn.ModuleList()

        self._dec_cross_attn_indices = []
        for i in reversed(range(len(channel_mults))):
            mult = channel_mults[i]
            out_ch = ch * mult
            skip_ch = skip_channels.pop()

            self.dec_blocks.append(ResBlock(cur_ch + skip_ch, out_ch, t_dim, dropout))
            cur_ch = out_ch

            dec_res = image_size // (2 ** i) if i < len(channel_mults) - 1 else cur_res
            if dec_res in attn_resolutions:
                self.dec_self_attns.append(SelfAttention(cur_ch, num_heads))
                self.dec_cross_attns.append(CrossAttention(cur_ch, num_heads))
                self._dec_cross_attn_indices.append(len(self.dec_cross_attns) - 1)
            else:
                self.dec_self_attns.append(nn.Identity())
                self.dec_cross_attns.append(nn.Identity())

            if i > 0:
                self.dec_ups.append(Upsample(cur_ch))
            else:
                self.dec_ups.append(nn.Identity())

        # ---- Output ----
        self.out_norm = nn.GroupNorm(min(32, cur_ch), cur_ch)
        self.out_conv = nn.Conv2d(cur_ch, in_channels, 3, padding=1)
        nn.init.zeros_(self.out_conv.weight)
        nn.init.zeros_(self.out_conv.bias)

        # Zero timestep embedding for y-trunk (doesn't depend on t)
        self.register_buffer("_zero_t_emb", torch.zeros(1, t_dim))

    def encode_y(self, y):
        """
        Run y-trunk encoder to get multi-resolution features.
        Returns list of features: [after_input_conv, after_level_0, after_level_1, ..., bottleneck]
        """
        zero_t = self._zero_t_emb.expand(y.shape[0], -1)
        h = self.y_input_conv(y)
        features = [h]

        for block, attn, down in zip(self.y_enc_blocks, self.y_enc_attns, self.y_enc_downs):
            h = block(h, zero_t)
            h = attn(h)
            features.append(h)
            h = down(h)

        # Bottleneck
        h = self.y_mid1(h, zero_t)
        h = self.y_mid_attn(h)
        h = self.y_mid2(h, zero_t)
        features.append(h)  # bottleneck features

        return features

    def forward(self, x_t, x_T, t, cond_mask=None, y_features=None):
        """
        Args:
            x_t: noisy intermediate [B, C, H, W]
            x_T: source conditioning image [B, C, H, W]
            t: timestep [B]
            cond_mask: [B, 1, 1, 1] binary mask. 0 = null conditioning (for CFG training).
            y_features: precomputed y-trunk features (optional, for efficient CFG inference).
                        If None, runs y-trunk on x_T.
        """
        # Get y-trunk features
        if y_features is None:
            y_features = self.encode_y(x_T)

        # Apply CFG conditioning dropout (zero out y features)
        if cond_mask is not None:
            y_features = [f * cond_mask for f in y_features]

        t_emb = self.time_mlp(t)

        # y_features layout: [input, level0, level1, ..., levelN, bottleneck]
        # index 0 = after input conv (for skip-level matching)
        # index 1..N = after each encoder level
        # index -1 = bottleneck

        # ---- x_t Encoder ----
        h = self.input_conv(x_t)
        skips = [h]

        for i, (block, self_attn, cross_attn, down) in enumerate(
            zip(self.enc_blocks, self.enc_self_attns, self.enc_cross_attns, self.enc_downs)
        ):
            h = block(h, t_emb)
            h = self_attn(h)
            # Cross-attend to y features at matching resolution
            if i in self._enc_cross_attn_indices:
                h = cross_attn(h, y_features[i + 1])  # +1 because features[0] is input conv
            skips.append(h)
            h = down(h)

        # ---- Bottleneck ----
        h = self.mid1(h, t_emb)
        h = self.mid_self_attn(h)
        h = self.mid_cross_attn(h, y_features[-1])  # bottleneck y features
        h = self.mid2(h, t_emb)

        # ---- Decoder ----
        for i, (block, self_attn, cross_attn, up) in enumerate(
            zip(self.dec_blocks, self.dec_self_attns, self.dec_cross_attns, self.dec_ups)
        ):
            h = torch.cat([h, skips.pop()], dim=1)
            h = block(h, t_emb)
            h = self_attn(h)
            if i in self._dec_cross_attn_indices:
                # Match to encoder-level y features (decoder mirrors encoder)
                # Decoder level i corresponds to encoder level (N-1-i)
                n_levels = len(self.enc_blocks)
                enc_idx = n_levels - i  # features index
                if enc_idx < len(y_features) - 1:  # -1 to exclude bottleneck
                    h = cross_attn(h, y_features[enc_idx])
            h = up(h)

        return self.out_conv(F.silu(self.out_norm(h)))


# ============================================================
# VQGAN wrapper for latent-space BBDM
# ============================================================

class LatentBBDM(nn.Module):
    """
    Wraps BBDMUNet with a frozen pretrained VQGAN for latent-space diffusion.

    Uses taming-transformers VQGAN-f4:
      - 256x256x3 images -> 64x64x3 pre-quant continuous latents
      - Latent features have ~unit variance (constrained by VQ codebook)
      - Compatible with SimpleUNet channel_mults=(1,2,4,4)

    The encode path uses encoder + quant_conv (pre-quantization features),
    NOT the quantized codes. This matches the original BBDM paper.
    """

    def __init__(
        self,
        unet_channels=128,
        latent_channels=3,
        image_size=256,
        latent_size=64,
        t_dim=256,
        attn_resolutions=(16, 32),
        channel_mults=(1, 2, 4, 4),
        num_heads=4,
        dropout=0.0,
        vqgan_ckpt=None,
        vqgan_config=None,
    ):
        super().__init__()
        self.image_size = image_size
        self.latent_size = latent_size
        self.latent_channels = latent_channels
        self.has_vqgan = False

        # UNet operates in latent space
        # For vq-f4 64x64x3 latents with (1,2,4,4):
        #   Stage 0: 64x64 x 128ch (attention at 32)
        #   Stage 1: 32x32 x 256ch (attention at 16)
        #   Stage 2: 16x16 x 512ch (attention at 16)
        #   Stage 3: 8x8 x 512ch (no downsample)
        #   Bottleneck: 8x8 x 512ch with attention
        self.unet = BBDMUNet(
            in_channels=latent_channels,
            base_channels=unet_channels,
            t_dim=t_dim,
            image_size=latent_size,
            channel_mults=channel_mults,
            attn_resolutions=attn_resolutions,
            num_heads=num_heads,
            dropout=dropout,
        )

        if vqgan_ckpt is not None:
            self._load_vqgan(vqgan_ckpt, vqgan_config)

    def _load_vqgan(self, ckpt_path, config_path):
        """Load taming-transformers VQGAN from checkpoint + config."""
        from omegaconf import OmegaConf
        from taming.models.vqgan import VQModel

        config = OmegaConf.load(config_path)
        params = config.model.params
        self.vqgan = VQModel(
            ddconfig=params.ddconfig,
            lossconfig=params.lossconfig,
            n_embed=params.n_embed,
            embed_dim=params.embed_dim,
        )

        sd = torch.load(ckpt_path, map_location="cpu", weights_only=False)
        if "state_dict" in sd:
            sd = sd["state_dict"]
        self.vqgan.load_state_dict(sd, strict=False)
        self.vqgan.eval()

        for p in self.vqgan.parameters():
            p.requires_grad_(False)

        self.has_vqgan = True
        self.latent_channels = params.ddconfig.z_channels
        print(f"Loaded VQGAN from {ckpt_path} (z_channels={self.latent_channels})")

    @torch.no_grad()
    def encode(self, x):
        """Encode to pre-quant continuous latent (NOT quantized codes)."""
        if not self.has_vqgan:
            return x
        h = self.vqgan.encoder(x)
        h = self.vqgan.quant_conv(h)
        return h

    @torch.no_grad()
    def decode(self, z):
        """Decode from pre-quant latent: quantize through VQ codebook, then decode."""
        if not self.has_vqgan:
            return z
        # Snap to nearest codebook entry (as in original BBDM paper)
        # This acts as free error correction for noisy predictions
        z_quant, _, _ = self.vqgan.quantize(z)
        return self.vqgan.decode(z_quant)

    def forward(self, z_t, z_T, t, cond_mask=None):
        return self.unet(z_t, z_T, t, cond_mask=cond_mask)


# ============================================================
# PatchGAN Discriminator
# ============================================================

class PatchGANDiscriminator(nn.Module):
    """
    PatchGAN discriminator (70x70 receptive field).

    Takes concatenated (source, target_or_generated) as input [B, 6, H, W].
    Outputs patch-level real/fake predictions [B, 1, N, N].
    """

    def __init__(self, in_channels=6, base_channels=64, n_layers=3):
        super().__init__()
        layers = [
            nn.Conv2d(in_channels, base_channels, 4, stride=2, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
        ]

        ch = base_channels
        for i in range(1, n_layers):
            out_ch = min(ch * 2, 512)
            layers += [
                nn.Conv2d(ch, out_ch, 4, stride=2, padding=1),
                nn.InstanceNorm2d(out_ch),
                nn.LeakyReLU(0.2, inplace=True),
            ]
            ch = out_ch

        # Second-to-last layer: stride 1
        out_ch = min(ch * 2, 512)
        layers += [
            nn.Conv2d(ch, out_ch, 4, stride=1, padding=1),
            nn.InstanceNorm2d(out_ch),
            nn.LeakyReLU(0.2, inplace=True),
        ]

        # Final layer: 1-channel prediction map
        layers += [nn.Conv2d(out_ch, 1, 4, stride=1, padding=1)]

        self.model = nn.Sequential(*layers)

    def forward(self, source, target):
        """
        Args:
            source: source domain image (HE) [B, 3, H, W]
            target: real or generated target image [B, 3, H, W]
        Returns:
            patch predictions [B, 1, N, N]
        """
        return self.model(torch.cat([source, target], dim=1))


# ============================================================
# Factory functions
# ============================================================

def SimpleUNet(in_channels=3, base_channels=64, t_dim=128, image_size=64,
               dual_output=False, use_x1_cond=True):
    """Drop-in replacement for old SimpleUNet. Now with attention.

    use_x1_cond=False removes x_1 from the U-Net input (3 channels in instead of 6).
    Conditioning then enters only via the bridge state x_t = α·x_0 + β·x_1 + γ·z.
    Useful as a cheat-blocker when schedule learning lets γ blow up.

    Set dual_output=True for the learned-schedule SNR-routed dual-target loss.
    """
    if image_size <= 64:
        attn_res = (16,)
        mults = (1, 2, 4)
    else:
        attn_res = (16, 32)
        mults = (1, 2, 4, 4)

    return BBDMUNet(
        in_channels=in_channels,
        base_channels=base_channels,
        t_dim=t_dim,
        image_size=image_size,
        channel_mults=mults,
        attn_resolutions=attn_res,
        dual_output=dual_output,
        use_x1_cond=use_x1_cond,
    )


def OriginalBBDMUNet(in_channels=3, base_channels=128, t_dim=256, image_size=64):
    """
    UNet matching the original BBDM paper (xuekt98/BBDM).

    Designed for latent-space diffusion on 64x64x3 VQGAN-f4 latents.
    Architecture: 64->32->16->8 bottleneck (3 stages, 2 downsamples)
    Channel mults: (1, 4, 8) -> [128, 512, 1024]
    Attention at 32, 16, and 8 — covers nearly the whole network.
    8 heads (64 dims/head at 512ch, matching LDM convention).
    """
    return BBDMUNet(
        in_channels=in_channels,
        base_channels=base_channels,
        t_dim=t_dim,
        image_size=image_size,
        channel_mults=(1, 4, 8),
        attn_resolutions=(32, 16, 8),
        num_heads=8,
        dropout=0.0,
    )


def DualTrunkUNet(in_channels=3, base_channels=128, t_dim=256, image_size=256):
    """Dual-trunk UNet with cross-attention conditioning from y-trunk."""
    if image_size <= 64:
        attn_res = (16,)
        mults = (1, 2, 4)
    else:
        attn_res = (16, 32)
        mults = (1, 2, 4, 4)

    return DualTrunkBBDMUNet(
        in_channels=in_channels,
        base_channels=base_channels,
        t_dim=t_dim,
        image_size=image_size,
        channel_mults=mults,
        attn_resolutions=attn_res,
    )


def DeepDualTrunkUNet(in_channels=3, base_channels=128, t_dim=256, image_size=256):
    """Deep dual-trunk UNet: 5 levels, cross-attention at 16x16 and 32x32."""
    return DualTrunkBBDMUNet(
        in_channels=in_channels,
        base_channels=base_channels,
        t_dim=t_dim,
        image_size=image_size,
        channel_mults=(1, 2, 4, 4, 4),
        attn_resolutions=(16, 32),
        num_heads=8,
        dropout=0.0,
    )


def DeepUNet(in_channels=3, base_channels=128, t_dim=256, image_size=256):
    """
    Deeper UNet that reaches 16x16 resolution with attention.

    Architecture: 256->128->64->32->16 (4 downsamples)
    Channel mults: (1, 2, 4, 4, 4) -> [128, 256, 512, 512, 512]
    Attention at 16x16 and 32x32 resolutions.
    Capped at 512 channels for fp16 stability.
    """
    return BBDMUNet(
        in_channels=in_channels,
        base_channels=base_channels,
        t_dim=t_dim,
        image_size=image_size,
        channel_mults=(1, 2, 4, 4, 4),
        attn_resolutions=(16, 32),
        num_heads=8,
        dropout=0.0,
    )