import math
from typing import Iterable, Tuple

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


def get_timestep_embedding(timesteps: torch.Tensor, embedding_dim: int) -> torch.Tensor:
    if timesteps.ndim != 1:
        raise ValueError("Timesteps should be a 1-D tensor")
    half_dim = embedding_dim // 2
    device = timesteps.device
    exponent = torch.arange(half_dim, dtype=torch.float32, device=device)
    exponent = -math.log(10000.0) * exponent / max(half_dim - 1, 1)
    freqs = torch.exp(exponent)
    args = timesteps.float().unsqueeze(1) * freqs.unsqueeze(0)
    emb = torch.cat([torch.sin(args), torch.cos(args)], dim=1)
    if embedding_dim % 2 == 1:
        emb = F.pad(emb, (0, 1))
    return emb


def nonlinearity(x: torch.Tensor) -> torch.Tensor:
    return x * torch.sigmoid(x)


def group_norm(num_channels: int) -> nn.GroupNorm:
    return nn.GroupNorm(num_groups=32, num_channels=num_channels, eps=1e-6, affine=True)


class Upsample(nn.Module):
    def __init__(self, in_channels: int, with_conv: bool) -> None:
        super().__init__()
        self.with_conv = with_conv
        if with_conv:
            self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)

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


class Downsample(nn.Module):
    def __init__(self, in_channels: int, with_conv: bool) -> None:
        super().__init__()
        self.with_conv = with_conv
        if with_conv:
            self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.with_conv:
            return self.conv(x)
        return F.avg_pool2d(x, kernel_size=2, stride=2)


class ResnetBlock(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        temb_channels: int,
        dropout: float,
    ) -> None:
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels

        self.norm1 = group_norm(in_channels)
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)

        self.temb_proj = nn.Sequential(
            nn.SiLU(),
            nn.Linear(temb_channels, out_channels),
        )

        self.norm2 = group_norm(out_channels)
        self.dropout = nn.Dropout(dropout)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)

        if in_channels != out_channels:
            self.shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1)
        else:
            self.shortcut = nn.Identity()

    def forward(self, x: torch.Tensor, temb: torch.Tensor) -> torch.Tensor:
        h = self.norm1(x)
        h = nonlinearity(h)
        h = self.conv1(h)

        h = h + self.temb_proj(temb)[:, :, None, None]

        h = self.norm2(h)
        h = nonlinearity(h)
        h = self.dropout(h)
        h = self.conv2(h)

        return self.shortcut(x) + h


class AttnBlock(nn.Module):
    def __init__(self, in_channels: int) -> None:
        super().__init__()
        self.in_channels = in_channels
        self.norm = group_norm(in_channels)
        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
        self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = self.norm(x)
        q = self.q(h)
        k = self.k(h)
        v = self.v(h)

        b, c, h_sp, w_sp = q.shape
        q = q.reshape(b, c, h_sp * w_sp)
        q = q.permute(0, 2, 1)
        k = k.reshape(b, c, h_sp * w_sp)
        w = torch.bmm(q, k)
        w = w * (c ** -0.5)
        w = torch.softmax(w, dim=2)

        v = v.reshape(b, c, h_sp * w_sp)
        w = w.permute(0, 2, 1)
        h = torch.bmm(v, w)
        h = h.reshape(b, c, h_sp, w_sp)

        h = self.proj_out(h)
        return x + h


def _to_tuple(value: Iterable[int] | int) -> Tuple[int, ...]:
    if isinstance(value, Iterable) and not isinstance(value, (str, bytes)):
        return tuple(value)
    return (int(value),)


class DiffusionUNet(nn.Module):
    def __init__(
        self,
        in_ch: int = 3,
        cond_ch: int = 3,
        out_ch: int = 3,
        base_ch: int = 128,
        channel_mults: Iterable[int] = (1, 1, 2, 2, 4, 4),
        num_res_blocks: int = 2,
        attn_resolutions: Iterable[int] = (16,),
        dropout: float = 0.0,
        resamp_with_conv: bool = True,
        num_variants: int = 3,
        image_size: int = 256,
    ) -> None:
        super().__init__()

        # [AttentionMask] Track diffusion vs conditioning channels (original line kept for rollback).
        # self.in_channels = in_ch + cond_ch
        self.x_ch = in_ch
        self.cond_ch = cond_ch
        if self.cond_ch <= 0:
            raise ValueError("cond_ch must be positive")
        self.in_channels = self.x_ch + self.cond_ch
        self.out_channels = out_ch
        self.ch = base_ch
        self.temb_ch = base_ch * 4
        self.num_resolutions = len(tuple(channel_mults))
        self.num_res_blocks = num_res_blocks
        self.image_size = image_size

        ch_mult = _to_tuple(channel_mults)
        attn_resolutions = _to_tuple(attn_resolutions)

        # Time embedding + optional variant embedding
        self.temb = nn.Module()
        self.temb.dense = nn.ModuleList(
            [
                nn.Linear(self.ch, self.temb_ch),
                nn.Linear(self.temb_ch, self.temb_ch),
            ]
        )
        self.variant_embed = nn.Embedding(num_variants, self.temb_ch) if num_variants > 0 else None

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

        curr_res = image_size
        in_ch_mult = (1,) + ch_mult
        self.down = nn.ModuleList()
        block_in = None
        for i_level in range(self.num_resolutions):
            block = nn.ModuleList()
            attn = nn.ModuleList()
            block_in = self.ch * in_ch_mult[i_level]
            block_out = self.ch * ch_mult[i_level]
            for _ in range(self.num_res_blocks):
                block.append(ResnetBlock(block_in, block_out, self.temb_ch, dropout))
                block_in = block_out
                if curr_res in attn_resolutions:
                    attn.append(AttnBlock(block_in))
            down = nn.Module()
            down.block = block
            down.attn = attn
            if i_level != self.num_resolutions - 1:
                down.downsample = Downsample(block_in, resamp_with_conv)
                curr_res = curr_res // 2
            self.down.append(down)

        self.mid = nn.Module()
        self.mid.block_1 = ResnetBlock(block_in, block_in, self.temb_ch, dropout)
        self.mid.attn_1 = AttnBlock(block_in)
        self.mid.block_2 = ResnetBlock(block_in, block_in, self.temb_ch, dropout)

        self.up = nn.ModuleList()
        for i_level in reversed(range(self.num_resolutions)):
            block = nn.ModuleList()
            attn = nn.ModuleList()
            block_out = self.ch * ch_mult[i_level]
            skip_in = self.ch * ch_mult[i_level]
            for i_block in range(self.num_res_blocks + 1):
                if i_block == self.num_res_blocks:
                    skip_in = self.ch * in_ch_mult[i_level]
                block.append(ResnetBlock(block_in + skip_in, block_out, self.temb_ch, dropout))
                block_in = block_out
                if curr_res in attn_resolutions:
                    attn.append(AttnBlock(block_in))
            up = nn.Module()
            up.block = block
            up.attn = attn
            if i_level != 0:
                up.upsample = Upsample(block_in, resamp_with_conv)
                curr_res = curr_res * 2
            self.up.insert(0, up)

        self.norm_out = group_norm(block_in)
        self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)

    def forward(
        self,
        x_t: torch.Tensor,
        t: torch.Tensor,
        y: torch.Tensor,
        variant_idx: torch.Tensor | None = None,
    ) -> torch.Tensor:
        # [AttentionMask] Validate channel counts before concatenation (original logic commented).
        # if y is not None:
        #     x = torch.cat([x_t, y], dim=1)
        # else:
        #     x = x_t
        if x_t.size(1) != self.x_ch:
            raise ValueError(f"Expected {self.x_ch} diffusion channels, got {x_t.size(1)}")
        if y is not None:
            if y.size(1) != self.cond_ch:
                raise ValueError(f"Expected {self.cond_ch} conditioning channels, got {y.size(1)}")
            x = torch.cat([x_t, y], dim=1)
        else:
            if self.cond_ch != 0:
                raise ValueError("cond_ch must be 0 when no conditioning tensor is provided")
            x = x_t

        temb = get_timestep_embedding(t, self.ch)
        temb = self.temb.dense[0](temb)
        temb = nonlinearity(temb)
        temb = self.temb.dense[1](temb)
        if self.variant_embed is not None and variant_idx is not None:
            temb = temb + self.variant_embed(variant_idx)

        hs = [self.conv_in(x)]
        h = hs[0]
        for i_level in range(self.num_resolutions):
            for i_block, res_block in enumerate(self.down[i_level].block):
                h = res_block(h, temb)
                if len(self.down[i_level].attn) > i_block:
                    h = self.down[i_level].attn[i_block](h)
                hs.append(h)
            if hasattr(self.down[i_level], "downsample"):
                h = self.down[i_level].downsample(h)
                hs.append(h)

        h = self.mid.block_1(h, temb)
        h = self.mid.attn_1(h)
        h = self.mid.block_2(h, temb)

        for i_level in reversed(range(self.num_resolutions)):
            for i_block, res_block in enumerate(self.up[i_level].block):
                h = torch.cat([h, hs.pop()], dim=1)
                h = res_block(h, temb)
                if len(self.up[i_level].attn) > i_block:
                    h = self.up[i_level].attn[i_block](h)
            if hasattr(self.up[i_level], "upsample"):
                h = self.up[i_level].upsample(h)

        h = self.norm_out(h)
        h = nonlinearity(h)
        return self.conv_out(h)


if __name__ == "__main__":
    model = DiffusionUNet(in_ch=3, cond_ch=3, base_ch=64, channel_mults=(1, 2, 2, 4), num_variants=3)
    model.eval()
    b = 2
    x_t = torch.randn(b, 3, 128, 128)
    y = torch.randn(b, 3, 128, 128)
    t = torch.randint(0, 1000, (b,))
    variants = torch.randint(0, 3, (b,))
    out = model(x_t, t, y, variants)
    print("Output shape:", out.shape)