import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange

class TimeEmbedding(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.proj = nn.Sequential(
            nn.Linear(1, dim),
            nn.SiLU(),
            nn.Linear(dim, dim)
        )
        
    def forward(self, t):
        return self.proj(t)

class Conv3DBlock(nn.Module):
    def __init__(self, in_ch, out_ch, time_dim):
        super().__init__()
        self.time_mlp = nn.Linear(time_dim, out_ch)
        self.conv = nn.Conv3d(in_ch, out_ch, kernel_size=3, padding=1)
        self.norm = nn.BatchNorm3d(out_ch)
        
    def forward(self, x, t_emb):
        t_emb = self.time_mlp(t_emb).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
        return F.silu(self.norm(self.conv(x) + t_emb))

class UNet3D(nn.Module):
    def __init__(self, in_ch=3, out_ch=3, text_dim=768):
        super().__init__()
        self.time_embed = TimeEmbedding(256)
        self.text_proj = nn.Linear(text_dim, 256)
        
        # Downsample
        self.down1 = Conv3DBlock(in_ch, 64, 256)
        self.down2 = Conv3DBlock(64, 128, 256)
        self.down3 = Conv3DBlock(128, 256, 256)
        
        # Upsample
        self.up1 = Conv3DBlock(256 + 128, 128, 256)
        self.up2 = Conv3DBlock(128 + 64, 64, 256)
        self.up3 = nn.Conv3d(64, out_ch, kernel_size=3, padding=1)
        
    def forward(self, x, t, text_emb):
        t_emb = self.time_embed(t)
        text_emb = self.text_proj(text_emb)
        c_emb = t_emb + text_emb
        
        # Downsample
        x1 = self.down1(x, c_emb)
        x2 = self.down2(F.max_pool3d(x1, 2), c_emb)
        x3 = self.down3(F.max_pool3d(x2, 2), c_emb)
        
        # Upsample
        x = F.interpolate(x3, scale_factor=2)
        x = self.up1(torch.cat([x, x2], dim=1), c_emb)
        x = F.interpolate(x, scale_factor=2)
        x = self.up2(torch.cat([x, x1], dim=1), c_emb)
        x = self.up3(x)
        return x