import torch.nn as nn
import math
import torch

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

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

        return emb

class ResBlock(nn.Module):
    def __init__(self, in_c, out_c, time_dim):
        super().__init__()

        self.conv1 = nn.Conv2d(
            in_c,
            out_c,
            3,
            padding=1
        )

        self.conv2 = nn.Conv2d(
            out_c,
            out_c,
            3,
            padding=1
        )

        self.time = nn.Linear(
            time_dim,
            out_c
        )

        self.norm1 = nn.GroupNorm(
            8,
            out_c
        )

        self.norm2 = nn.GroupNorm(
            8,
            out_c
        )

        self.act = nn.SiLU()


        self.skip = (
            nn.Conv2d(in_c,out_c,1)
            if in_c != out_c
            else nn.Identity()
        )


    def forward(self, x, t):
          h = self.conv1(x)
          h = self.norm1(h)
          h = self.act(h)

          t = self.time(t)
          t = self.act(t)

          h = h + t[:, :, None, None]

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

          return self.act(h + self.skip(x))

class SelfAttention(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.channels = channels
        self.mha = nn.MultiheadAttention(embed_dim=channels, num_heads=4, batch_first=True)
        self.ln = nn.LayerNorm([channels])
        self.ff_net = nn.Sequential(
            nn.Linear(channels, channels),
            nn.SiLU(),
            nn.Linear(channels, channels),
        )

    def forward(self, x):
        B, C, H, W = x.shape
        
        x_flat = x.view(B, C, H * W).transpose(1, 2)
        x_norm = self.ln(x_flat)
        attn_out, _ = self.mha(x_norm, x_norm, x_norm)
        attn_out = attn_out + x_flat
        ff_out = self.ff_net(self.ln(attn_out)) + attn_out
        
        return ff_out.transpose(1, 2).view(B, C, H, W)


class UNetWithAttention(nn.Module):
    def __init__(self, time_dim=256):
        super().__init__()

        self.time = nn.Sequential(
            TimeEmbedding(time_dim),
            nn.Linear(time_dim, time_dim),
            nn.SiLU()
        )

        self.input = nn.Conv2d(3, 64, 3, padding=1)

        self.down1 = ResBlock(64, 128, time_dim)
        self.pool1 = nn.Conv2d(128, 128, 4, stride=2, padding=1)

        self.down2 = ResBlock(128, 256, time_dim)
        self.down2_attn = SelfAttention(256)

        self.mid = ResBlock(256, 256, time_dim)
        self.mid_attn = SelfAttention(256)

        self.upblock1 = ResBlock(512, 256, time_dim)
        self.upblock1_attn = SelfAttention(256)

        self.up1 = nn.ConvTranspose2d(256, 128, 4, stride=2, padding=1) # Up to 64x64
        self.upblock2 = ResBlock(256, 128, time_dim)

        self.output = nn.Sequential(
            nn.Conv2d(128, 64, 3, padding=1),
            nn.SiLU(),
            nn.Conv2d(64, 3, 1)
        )

    def forward(self, x, t):
        t_embed = self.time(t)
        x = self.input(x)

        x1 = self.down1(x, t_embed)
        x = self.pool1(x1)

        x2 = self.down2(x, t_embed)
        x2 = self.down2_attn(x2)

        x = self.mid(x2, t_embed)
        x = self.mid_attn(x)

        x = torch.cat([x, x2], dim=1)
        x = self.upblock1(x, t_embed)
        x = self.upblock1_attn(x)

        x = self.up1(x)
        x = torch.cat([x, x1], dim=1)
        x = self.upblock2(x, t_embed)

        return self.output(x)