"""Time-conditioned U-Net for DDPM (Phase 4).

Architecture follows Ho et al. (2020) with options from Nichol & Dhariwal (2021):
  - Sinusoidal time embedding → MLP → added to every ResBlock
  - GroupNorm (32 groups) + SiLU activations throughout
  - Self-attention at configurable spatial resolutions
  - Upsample(nearest) + Conv in the decoder — no checkerboard artefacts

Registered as kind="ddpm".
"""
import math

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

from src.models import register

_GN = 32  # GroupNorm groups — all channel counts used here are multiples of 32


# ── Time embedding ────────────────────────────────────────────────────────────

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

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        half = self.dim // 2
        freqs = torch.exp(
            -math.log(10000) * torch.arange(half, device=t.device, dtype=torch.float) / half
        )
        angles = t[:, None].float() * freqs[None]  # (B, half)
        return torch.cat([angles.sin(), angles.cos()], dim=-1)  # (B, dim)


# ── Core building blocks ──────────────────────────────────────────────────────

class ResBlock(nn.Module):
    """ResNet block with time-embedding injection (additive, after first conv)."""

    def __init__(self, in_ch: int, out_ch: int, t_emb_dim: int, dropout: float = 0.1):
        super().__init__()
        self.norm1   = nn.GroupNorm(_GN, in_ch)
        self.conv1   = nn.Conv2d(in_ch, out_ch, 3, padding=1)
        self.t_proj  = nn.Linear(t_emb_dim, out_ch)
        self.norm2   = nn.GroupNorm(_GN, out_ch)
        self.drop    = nn.Dropout(dropout)
        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: torch.Tensor, t_emb: torch.Tensor) -> torch.Tensor:
        h = self.conv1(F.silu(self.norm1(x)))
        h = h + self.t_proj(F.silu(t_emb))[:, :, None, None]
        h = self.conv2(self.drop(F.silu(self.norm2(h))))
        return h + self.skip(x)


class AttentionBlock(nn.Module):
    """Single-head self-attention with GroupNorm pre-norm and residual."""

    def __init__(self, ch: int):
        super().__init__()
        self.norm = nn.GroupNorm(_GN, ch)
        self.qkv  = nn.Conv2d(ch, ch * 3, 1, bias=False)
        self.proj = nn.Conv2d(ch, ch, 1)
        self._scale = ch ** -0.5

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b, c, h, w = x.shape
        n = h * w
        qkv = self.norm(x)
        q, k, v = self.qkv(qkv).reshape(b, 3, c, n).unbind(1)   # each (b, c, n)
        attn = torch.softmax(q.transpose(-2, -1) @ k * self._scale, dim=-1)  # (b, n, n)
        out  = (v @ attn.transpose(-2, -1)).reshape(b, c, h, w)
        return x + self.proj(out)


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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.conv(x)


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

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


# ── Down / Up blocks ──────────────────────────────────────────────────────────

class DownBlock(nn.Module):
    def __init__(
        self,
        in_ch: int,
        out_ch: int,
        t_emb_dim: int,
        num_res_blocks: int,
        with_attn: bool,
        dropout: float,
    ):
        super().__init__()
        self.resnets = nn.ModuleList(
            ResBlock(in_ch if j == 0 else out_ch, out_ch, t_emb_dim, dropout)
            for j in range(num_res_blocks)
        )
        self.attn = AttentionBlock(out_ch) if with_attn else nn.Identity()

    def forward(self, x: torch.Tensor, t_emb: torch.Tensor) -> torch.Tensor:
        for res in self.resnets:
            x = res(x, t_emb)
        return self.attn(x)


class UpBlock(nn.Module):
    def __init__(
        self,
        in_ch: int,
        skip_ch: int,
        out_ch: int,
        t_emb_dim: int,
        num_res_blocks: int,
        with_attn: bool,
        dropout: float,
    ):
        super().__init__()
        # First ResBlock absorbs the skip-connection channels via concat
        self.resnets = nn.ModuleList(
            ResBlock(
                (in_ch + skip_ch) if j == 0 else out_ch,
                out_ch,
                t_emb_dim,
                dropout,
            )
            for j in range(num_res_blocks + 1)  # +1 to consume the concat
        )
        self.attn = AttentionBlock(out_ch) if with_attn else nn.Identity()

    def forward(
        self, x: torch.Tensor, skip: torch.Tensor, t_emb: torch.Tensor
    ) -> torch.Tensor:
        x = torch.cat([x, skip], dim=1)
        for res in self.resnets:
            x = res(x, t_emb)
        return self.attn(x)


# ── U-Net ─────────────────────────────────────────────────────────────────────

class UNet(nn.Module):
    """Time-conditioned U-Net.

    image_size    — must be a power-of-two; 64 or 128 recommended.
    base_ch       — base channel count (128 for phases 4.1–4.3, 192 for 4.4).
    ch_mult       — channel multipliers per resolution level.
    attn_resolutions — spatial resolutions at which attention is inserted.
    num_res_blocks — ResBlocks per level (in both down and up paths).
    dropout       — applied inside every ResBlock.
    """

    def __init__(
        self,
        image_size: int = 64,
        base_ch: int = 128,
        ch_mult: tuple = (1, 2, 2, 2),
        attn_resolutions: tuple = (16, 8),
        num_res_blocks: int = 2,
        dropout: float = 0.1,
    ):
        super().__init__()
        n_levels  = len(ch_mult)
        chs       = [base_ch * m for m in ch_mult]
        t_emb_dim = base_ch * 4

        # ── Time embedding ────────────────────────────────────────────────
        self.time_embed = nn.Sequential(
            SinusoidalPosEmb(base_ch),
            nn.Linear(base_ch, t_emb_dim),
            nn.SiLU(),
            nn.Linear(t_emb_dim, t_emb_dim),
        )

        # ── Input projection ──────────────────────────────────────────────
        self.in_conv = nn.Conv2d(3, chs[0], 3, padding=1)

        # ── Down path ─────────────────────────────────────────────────────
        self.down_blocks  = nn.ModuleList()
        self.downsamples  = nn.ModuleList()
        cur_res = image_size
        prev_ch = chs[0]
        for i, ch in enumerate(chs):
            with_attn = (cur_res in attn_resolutions)
            self.down_blocks.append(
                DownBlock(prev_ch, ch, t_emb_dim, num_res_blocks, with_attn, dropout)
            )
            if i < n_levels - 1:
                self.downsamples.append(Downsample(ch))
                cur_res //= 2
            prev_ch = ch

        # ── Middle ────────────────────────────────────────────────────────
        mid_ch = chs[-1]
        self.mid_res1 = ResBlock(mid_ch, mid_ch, t_emb_dim, dropout)
        self.mid_attn = AttentionBlock(mid_ch)
        self.mid_res2 = ResBlock(mid_ch, mid_ch, t_emb_dim, dropout)

        # ── Up path ───────────────────────────────────────────────────────
        # Mirrors down path: iterate chs in reverse; skip_ch = chs[n_levels-1-i].
        self.up_blocks  = nn.ModuleList()
        self.upsamples  = nn.ModuleList()
        in_ch = mid_ch
        for i in range(n_levels):
            level     = n_levels - 1 - i      # index from deep (n-1) to shallow (0)
            skip_ch   = chs[level]
            out_ch    = chs[level - 1] if level > 0 else chs[0]
            with_attn = (cur_res in attn_resolutions)
            self.up_blocks.append(
                UpBlock(in_ch, skip_ch, out_ch, t_emb_dim, num_res_blocks, with_attn, dropout)
            )
            if level > 0:
                self.upsamples.append(Upsample(out_ch))
                cur_res *= 2
            in_ch = out_ch

        # ── Output ────────────────────────────────────────────────────────
        self.out_norm = nn.GroupNorm(_GN, chs[0])
        self.out_conv = nn.Conv2d(chs[0], 3, 3, padding=1)

    def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        t_emb = self.time_embed(t)

        x = self.in_conv(x)

        # Down
        skips = []
        ds_idx = 0
        for i, block in enumerate(self.down_blocks):
            x = block(x, t_emb)
            skips.append(x)
            if ds_idx < len(self.downsamples):
                x = self.downsamples[ds_idx](x)
                ds_idx += 1

        # Middle
        x = self.mid_res1(x, t_emb)
        x = self.mid_attn(x)
        x = self.mid_res2(x, t_emb)

        # Up
        us_idx = 0
        for i, block in enumerate(self.up_blocks):
            x = block(x, skips[-(i + 1)], t_emb)
            if us_idx < len(self.upsamples):
                x = self.upsamples[us_idx](x)
                us_idx += 1

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


def _build(cfg: dict):
    return UNet(
        image_size       = cfg.get("image_size", 64),
        base_ch          = cfg.get("base_ch", 128),
        ch_mult          = tuple(cfg.get("ch_mult", [1, 2, 2, 2])),
        attn_resolutions = tuple(cfg.get("attn_resolutions", [16, 8])),
        num_res_blocks   = cfg.get("num_res_blocks", 2),
        dropout          = cfg.get("dropout", 0.1),
    )


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