"""
Conditional U-Net Architecture for Diffusion Model
"""

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


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

    def forward(self, time):
        device = time.device
        half_dim = self.dim // 2
        embeddings = math.log(10000) / (half_dim - 1)
        embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)
        embeddings = time[:, None] * embeddings[None, :]
        return torch.cat([torch.sin(embeddings), torch.cos(embeddings)], dim=-1)


class LabelEmbedding(nn.Module):
    def __init__(self, label_dim, emb_dim):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(label_dim, emb_dim),
            nn.SiLU(),
            nn.Linear(emb_dim, emb_dim)
        )

    def forward(self, labels):
        return self.mlp(labels)


class ResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, time_emb_dim, dropout=0.1):
        super().__init__()
        self.conv1 = nn.Sequential(
            nn.GroupNorm(8, in_channels),
            nn.SiLU(),
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
        )
        self.time_emb = nn.Sequential(nn.SiLU(), nn.Linear(time_emb_dim, out_channels))
        self.conv2 = nn.Sequential(
            nn.GroupNorm(8, out_channels),
            nn.SiLU(),
            nn.Dropout(dropout),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
        )
        self.shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1) if in_channels != out_channels else nn.Identity()

    def forward(self, x, emb):
        h = self.conv1(x)
        h = h + self.time_emb(emb)[:, :, None, None]
        h = self.conv2(h)
        return h + self.shortcut(x)


class AttentionBlock(nn.Module):
    def __init__(self, channels, num_heads=4):
        super().__init__()
        self.channels = channels
        self.num_heads = num_heads
        self.norm = nn.GroupNorm(8, channels)
        self.qkv = nn.Conv2d(channels, channels * 3, kernel_size=1)
        self.proj = nn.Conv2d(channels, channels, kernel_size=1)

    def forward(self, x):
        B, C, H, W = x.shape
        h = self.norm(x)
        q, k, v = self.qkv(h).chunk(3, dim=1)
        head_dim = C // self.num_heads
        q = q.view(B, self.num_heads, head_dim, H*W).transpose(2, 3)
        k = k.view(B, self.num_heads, head_dim, H*W).transpose(2, 3)
        v = v.view(B, self.num_heads, head_dim, H*W).transpose(2, 3)

        h = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0)
        h = h.transpose(2, 3).reshape(B, C, H, W)
        return x + self.proj(h)


class ConditionalUNet(nn.Module):
    def __init__(self, in_channels=1, out_channels=1, label_dim=2,
                 base_channels=64, channel_multipliers=(1,2,4,8),
                 attention_levels=(2,3), dropout=0.1, time_emb_dim=256, label_emb_dim=256):
        super().__init__()
        self.label_dim = label_dim

        self.time_embedding = TimeEmbedding(time_emb_dim)
        self.time_mlp = nn.Sequential(nn.Linear(time_emb_dim, time_emb_dim*4), nn.SiLU(), nn.Linear(time_emb_dim*4, time_emb_dim))

        self.label_embedding = LabelEmbedding(label_dim, label_emb_dim)
        self.combined_emb_dim = time_emb_dim + label_emb_dim
        self.combined_mlp = nn.Sequential(nn.Linear(self.combined_emb_dim, time_emb_dim*4), nn.SiLU(), nn.Linear(time_emb_dim*4, time_emb_dim))

        self.conv_in = nn.Conv2d(in_channels, base_channels, kernel_size=3, padding=1)

        self.down_blocks = nn.ModuleList()
        channels = [base_channels]
        now_channels = base_channels

        for i, mult in enumerate(channel_multipliers):
            out_ch = base_channels * mult
            for _ in range(2):
                self.down_blocks.append(ResidualBlock(now_channels, out_ch, time_emb_dim, dropout))
                if i in attention_levels:
                    self.down_blocks.append(AttentionBlock(out_ch))
                now_channels = out_ch
                channels.append(now_channels)
            if i != len(channel_multipliers) - 1:
                self.down_blocks.append(nn.Conv2d(now_channels, now_channels, kernel_size=3, stride=2, padding=1))
                channels.append(now_channels)

        self.middle = nn.ModuleList([
            ResidualBlock(now_channels, now_channels, time_emb_dim, dropout),
            AttentionBlock(now_channels),
            ResidualBlock(now_channels, now_channels, time_emb_dim, dropout)
        ])

        self.up_blocks = nn.ModuleList()
        for i, mult in reversed(list(enumerate(channel_multipliers))):
            out_ch = base_channels * mult
            for _ in range(3):
                self.up_blocks.append(ResidualBlock(now_channels + channels.pop(), out_ch, time_emb_dim, dropout))
                if i in attention_levels:
                    self.up_blocks.append(AttentionBlock(out_ch))
                now_channels = out_ch
            if i != 0:
                self.up_blocks.append(nn.ConvTranspose2d(now_channels, now_channels, kernel_size=4, stride=2, padding=1))

        self.conv_out = nn.Sequential(
            nn.GroupNorm(8, now_channels),
            nn.SiLU(),
            nn.Conv2d(now_channels, out_channels, kernel_size=3, padding=1)
        )

    def forward(self, x, t, labels=None):
        t_emb = self.time_embedding(t)
        t_emb = self.time_mlp(t_emb)

        if labels is not None:
            label_emb = self.label_embedding(labels)
            combined = torch.cat([t_emb, label_emb], dim=-1)
            emb = self.combined_mlp(combined)
        else:
            emb = t_emb

        h = self.conv_in(x)
        skips = [h]

        for module in self.down_blocks:
            if isinstance(module, ResidualBlock):
                h = module(h, emb)
                skips.append(h)
            elif isinstance(module, AttentionBlock):
                h = module(h)
            else:
                h = module(h)
                skips.append(h)

        for module in self.middle:
            if isinstance(module, ResidualBlock):
                h = module(h, emb)
            else:
                h = module(h)

        for module in self.up_blocks:
            if isinstance(module, ResidualBlock):
                h = torch.cat([h, skips.pop()], dim=1)
                h = module(h, emb)
            elif isinstance(module, AttentionBlock):
                h = module(h)
            else:
                h = module(h)

        return self.conv_out(h)