import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import CLIPTextModel, CLIPTokenizer

class TimeEmbedding(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim
        half_dim = dim // 2
        emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb)
        self.register_buffer('emb', emb)
    
    def forward(self, time):
        emb = time[:, None] * self.emb[None, :]
        emb = torch.cat((torch.sin(emb), torch.cos(emb)), dim=-1)
        return emb

class AttentionBlock(nn.Module):
    def __init__(self, channels, num_heads=4):
        super().__init__()
        self.num_heads = num_heads
        self.scale = (channels // num_heads) ** -0.5
        
        self.norm = nn.GroupNorm(32, channels)
        self.qkv = nn.Conv2d(channels, channels * 3, 1)
        self.proj = nn.Conv2d(channels, channels, 1)
    
    def forward(self, x):
        b, c, h, w = x.shape
        qkv = self.qkv(self.norm(x))
        q, k, v = qkv.chunk(3, dim=1)
        
        q = q.view(b, self.num_heads, -1, h * w).permute(0, 1, 3, 2)
        k = k.view(b, self.num_heads, -1, h * w)
        v = v.view(b, self.num_heads, -1, h * w)
        
        attn = torch.softmax((q @ k) * self.scale, dim=-1)
        x = (attn @ v).permute(0, 1, 3, 2).reshape(b, -1, h, w)
        return self.proj(x) + x

class ResBlock(nn.Module):
    def __init__(self, in_ch, out_ch, time_emb_dim, text_emb_dim, dropout=0.1):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.SiLU(),
            nn.Linear(time_emb_dim + text_emb_dim, out_ch * 2)
        
        self.block1 = nn.Sequential(
            nn.GroupNorm(32, in_ch),
            nn.SiLU(),
            nn.Conv2d(in_ch, out_ch, 3, padding=1))
        
        self.block2 = nn.Sequential(
            nn.GroupNorm(32, out_ch),
            nn.SiLU(),
            nn.Dropout(dropout),
            nn.Conv2d(out_ch, out_ch, 3, padding=1))
        
        self.res_conv = nn.Conv2d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity()
    
    def forward(self, x, time_emb, text_emb):
        emb = self.mlp(torch.cat([time_emb, text_emb], dim=-1))
        scale, shift = torch.chunk(emb, 2, dim=1)
        
        h = self.block1(x)
        h = h * (1 + scale[:, :, None, None]) + shift[:, :, None, None]
        h = self.block2(h)
        
        return h + self.res_conv(x)

class UNet(nn.Module):
    def __init__(self, in_channels=3, out_channels=3, dim=64, dim_mults=(1, 2, 4, 8)):
        super().__init__()
        dims = [dim * m for m in dim_mults]
        in_out = list(zip(dims[:-1], dims[1:]))
        
        # Time and text embeddings
        self.time_mlp = nn.Sequential(
            TimeEmbedding(dim),
            nn.Linear(dim, dim * 4),
            nn.SiLU(),
            nn.Linear(dim * 4, dim))
        
        # Text conditioning
        self.text_proj = nn.Linear(768, dim * 4)
        
        # Initial convolution
        self.init_conv = nn.Conv2d(in_channels, dim, 3, padding=1)
        
        # Downsample blocks
        self.downs = nn.ModuleList()
        for ind, (in_dim, out_dim) in enumerate(in_out):
            is_last = ind >= (len(in_out) - 1)
            self.downs.append(nn.ModuleList([
                ResBlock(in_dim, in_dim, dim, dim * 4),
                ResBlock(in_dim, in_dim, dim, dim * 4),
                AttentionBlock(in_dim),
                nn.Conv2d(in_dim, out_dim, 3, stride=2, padding=1) if not is_last else nn.Conv2d(in_dim, out_dim, 3, padding=1)
            ]))
        
        # Middle blocks
        self.mid_block1 = ResBlock(dims[-1], dims[-1], dim, dim * 4)
        self.mid_attn = AttentionBlock(dims[-1])
        self.mid_block2 = ResBlock(dims[-1], dims[-1], dim, dim * 4)
        
        # Upsample blocks
        self.ups = nn.ModuleList()
        for ind, (in_dim, out_dim) in enumerate(reversed(in_out)):
            is_last = ind >= (len(in_out) - 1)
            self.ups.append(nn.ModuleList([
                ResBlock(out_dim + in_dim, out_dim, dim, dim * 4),
                ResBlock(out_dim + in_dim, out_dim, dim, dim * 4),
                AttentionBlock(out_dim),
                nn.ConvTranspose2d(out_dim, out_dim, 4, 2, 1) if not is_last else nn.Identity()
            ]))
        
        # Final blocks
        self.final_block1 = ResBlock(dim * 2, dim, dim, dim * 4)
        self.final_block2 = ResBlock(dim, dim, dim, dim * 4)
        self.final_conv = nn.Conv2d(dim, out_channels, 3, padding=1)
    
    def forward(self, x, time, text_emb):
        t = self.time_mlp(time)
        text_emb = self.text_proj(text_emb)
        
        x = self.init_conv(x)
        h = [x]
        
        # Downsample
        for block1, block2, attn, downsample in self.downs:
            x = block1(x, t, text_emb)
            x = block2(x, t, text_emb)
            x = attn(x)
            h.append(x)
            x = downsample(x)
        
        # Bottleneck
        x = self.mid_block1(x, t, text_emb)
        x = self.mid_attn(x)
        x = self.mid_block2(x, t, text_emb)
        
        # Upsample
        for block1, block2, attn, upsample in self.ups:
            x = torch.cat([x, h.pop()], dim=1)
            x = block1(x, t, text_emb)
            x = block2(x, t, text_emb)
            x = attn(x)
            x = upsample(x)
        
        # Final
        x = torch.cat([x, h.pop()], dim=1)
        x = self.final_block1(x, t, text_emb)
        x = self.final_block2(x, t, text_emb)
        return self.final_conv(x)

class DiffusionModel(nn.Module):
    def __init__(self, model, betas, device):
        super().__init__()
        self.model = model
        self.betas = betas
        self.alphas = 1. - betas
        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
        self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)
        self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1. - self.alphas_cumprod)
        self.device = device
        
        # CLIP model for text conditioning
        self.clip = CLIPTextModel.from_pretrained("openai/clip-vit-base-patch32")
        self.tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
        for param in self.clip.parameters():
            param.requires_grad = False
    
    def get_text_emb(self, prompts):
        inputs = self.tokenizer(prompts, padding=True, return_tensors="pt").to(self.device)
        return self.clip(**inputs).last_hidden_state.mean(dim=1)
    
    def q_sample(self, x_start, t, noise=None):
        if noise is None:
            noise = torch.randn_like(x_start)
        
        sqrt_alpha_cumprod = self.sqrt_alphas_cumprod[t].view(-1, 1, 1, 1)
        sqrt_one_minus_alpha_cumprod = self.sqrt_one_minus_alphas_cumprod[t].view(-1, 1, 1, 1)
        
        return sqrt_alpha_cumprod * x_start + sqrt_one_minus_alpha_cumprod * noise
    
    def p_losses(self, x_start, text, t, noise=None):
        if noise is None:
            noise = torch.randn_like(x_start)
        
        x_noisy = self.q_sample(x_start, t, noise)
        text_emb = self.get_text_emb(text)
        predicted_noise = self.model(x_noisy, t, text_emb)
        
        return F.mse_loss(noise, predicted_noise)
    
    @torch.no_grad()
    def sample(self, prompts, image_size=256, batch_size=4, channels=3, cfg_scale=7.5):
        shape = (batch_size, channels, image_size, image_size)
        x = torch.randn(shape, device=self.device)
        
        text_emb = self.get_text_emb(prompts)
        uncond_emb = self.get_text_emb([""] * batch_size)
        
        for i in reversed(range(0, len(self.betas))):
            t = torch.full((batch_size,), i, device=self.device, dtype=torch.long)
            
            # Classifier-free guidance
            noise_pred = self.model(x, t, text_emb)
            noise_pred_uncond = self.model(x, t, uncond_emb)
            noise_pred = noise_pred_uncond + cfg_scale * (noise_pred - noise_pred_uncond)
            
            alpha = self.alphas[t].view(-1, 1, 1, 1)
            alpha_cumprod = self.alphas_cumprod[t].view(-1, 1, 1, 1)
            beta = self.betas[t].view(-1, 1, 1, 1)
            
            if i > 0:
                noise = torch.randn_like(x)
            else:
                noise = torch.zeros_like(x)
            
            x = (1 / torch.sqrt(alpha)) * (x - ((1 - alpha) / torch.sqrt(1 - alpha_cumprod)) * noise_pred) + torch.sqrt(beta) * noise
        
        x = (x.clamp(-1, 1) + 1) / 2
        x = (x * 255).type(torch.uint8)
        return x