import torch 
import torch.nn as nn 
import math 

#Varicence scheduler 
T = 1000
beta_start = 1e-4
beta_end = 0.02
betas = torch.linspace(beta_start, beta_end, steps=T)
alphas = 1. - betas 
alphas_cumprod = torch.cumprod(alphas, dim=0)

#SinusoidalPositionEmbedding
class Sinusoidal_embedding(torch.nn.Module):
    def __init__(self, dim:int):
        super().__init__()
        assert dim % 2 == 0 , 'Embeddings must be divisble by 2'
        self.dim = dim 

    def forward(self, time_stamps:torch.Tensor) -> torch.Tensor :
        half_dim = self.dim // 2
        scale = math.log(10000)/(half_dim -1 )
        freq = torch.exp(torch.arange(half_dim, dtype=torch.float32, device=time_stamps.device)* -scale)
        embeddings = time_stamps[:, None] * freq[None, :]
        embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)
        return embeddings

#Time Embeddings 
class Time_embeddings(torch.nn.Module):
    def __init__(self, in_dim:int, out_dim:int):
        super().__init__()
        self.Sinusoidal_waves = Sinusoidal_embedding(in_dim)

        self.mlp = nn.Sequential(
            nn.Linear(in_dim, out_dim),
            nn.SiLU(),
            nn.Linear(out_dim , out_dim)
        )

    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
        time_emb = self.Sinusoidal_waves(timesteps)
        return self.mlp(time_emb)

#Resblock 
class resblock(torch.nn.Module):
    def __init__(self, in_channels, out_channels, time_emb_dim):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
        self.norm1 = nn.GroupNorm(8, out_channels)
        self.norm2 = nn.GroupNorm(8, out_channels)
        self.silu = nn.SiLU()
        self.time_mlp = nn.Linear(time_emb_dim, out_channels)
        #Matching in channels and out channels for skip connection 
        self.residual_conv = nn.Conv2d(in_channels, out_channels, 1) if in_channels != out_channels else nn.Identity()

    def forward(self, x, time_dim):
        h = self.norm1(self.conv1(x))
        h = self.silu(h)
        #Tensor broadcasting (b, out_channels) -> (b, out_channels, 1, 1)
        time_proj = self.silu(self.time_mlp(time_dim))[:, :, None, None]
        h = h + time_proj
        h = self.norm2(self.conv2(h))
        h = self.silu(h)
        return h + self.residual_conv(x)

#Downsmaple and upsmaple
class down_sample(torch.nn.Module):
    def __init__(self, channels):
        super().__ini__()
        self.conv = nn.Conv2d(channels, channels, 4, stride=2, padding=1)
    def forward(self, x):
        return self.conv(x)

class up_sample(torch.nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv = nn.ConvTranspose2d(channels, channels, 4, stride=2, padding=1)
    def forward(self, x):
        return self.conv(x)

#Assembling Unet class 
class Unet(torch.nn.Module):
    def __init__(self, in_channels=3, time_dim=128):
        super().__init__()
        #Time embedding diemension 
        self.time_emb = Time_embeddings(in_dim=time_dim, out_dim=256)
        self.inc = nn.Conv2d(3, 64, 3, padding=1)

        #Encoder class 
        self.resblock1 = resblock(64, 128, time_emb_dim=256)
        self.down1 = down_sample(128)

        self.resblock2 = resblock(128, 256, time_emb_dim=256)
        self.down2 = down_sample(256)

        #Bottleneck 
        self.bot1 = resblock(256,256, time_emb_dim=256)
        self.bot2 = resblock(256,256, time_emb_dim=256)

        #Decoder class
        self.up1 = up_sample(256)
        self.resblock3 = resblock(512, 128, time_emb_dim=256)

        self.up2 = up_sample(128)
        self.resblock4 = resblock(256, 64, time_emb_dim=256)

        #Final output layer 
        self.outc = nn.Conv2d(64, in_channels, 3, padding=1)

    def forward(self, x, t):

        t_emb = self.time_emb(t)
        x1 = self.inc(x)

        x2 = self.resblock1(x1, t_emb)
        x3 = self.down1(x2)
        x4 = self.resblock2(x3, t_emb)
        x5 = self.down2(x4)

        #bottle neck 
        x5 = self.bot1(x5, t_emb)
        x5 = self.bot2(x5, t_emb)

        #decoder
        x6 = self.up1(x5)
        x6 = torch.cat([x4, x6], dim=1)
        x6 = self.resblock3(x6, t_emb)

        x7 = self.up2(x6)
        x7 = torch.cat([x2, x7], dim=1)
        x7 = self.resblock4(x7, t_emb)

        return self.outc(x7)

#EMA weights 
class EMA():
    def __init__(self, model, decay=0.9999):
        self.model = model
        self.decay = decay 
        self.shadow = {}
        self.backup = {}

    def register(self):
        for name , param in self.model.named_parameters():
            if param.requires_grad:
                self.shadow[name] = param.data.clone()

    def upadate(self):
        for name, param in self.model.named_parameters():
            if param.requires_gard:
                assert name in self.shadow
                new_avg = self.decay*self.shadow[name] + ((1. - self.decay)*param)
                self.shadow[name] = new_avg.clone()

    def apply_shadow(self):
        for name , param in self.model.named_parameters():
            if param.requires_grad:
                assert name in self.shadow
                self.backup = param.data.clone()
                param.data.copy_(self.shadow[name])

    def restore(self):
        for name, param in self.model.named_parameters():
            if param.require_grad:
                assert name in self.backup 
                param.data.copy_(self.backup[name])
                self.backup = {}

#DDIM 
@torch.no_grad()
def sample_ddim(model, T,img_size,  batch_size=1, channels=3,ddim_step=50, device='cpu'):
    model.eval()
    alphas_cumprod_device = alphas_cumprod.to(device)
    timestep = torch.linspace(0, T-1, steps=ddim_step, dtype=torch.long, device=device)
    img = torch.randn(batch_size, channels, img_size, img_size).to(device)
    
    for i in reversed(range(len(timestep))):
        t_current =timestep[i]
        t_preq = timestep[i-1] if i > 0 else -1
        t_tensor = torch.full((batch_size,), t_current, dtype=torch.long, device=device)
        
        pred_noise = model(img , t_tensor)
        alphas_bar_t = alphas_cumprod_device[t_current]
        alphas_bar_t_preq = alphas_cumprod_device[t_preq] if t_preq >= 0 else torch.tensor(1.0, device=device)
        pred_x0 = (img - torch.sqrt(1 - alphas_bar_t) * pred_noise)/(torch.sqrt(alphas_bar_t))
        pred_x0 = torch.clamp(pred_x0, -1.0, 1.0)
        pred_dir = torch.sqrt(1 - alphas_bar_t_preq) * pred_noise
        img = torch.sqrt(alphas_bar_t_preq) * pred_x0 + pred_dir
        
    return img