import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms


# ---- 1. Noise schedule ----

def linear_beta_schedule(T: int, beta_start=1e-4, beta_end=0.02):
    return torch.linspace(beta_start, beta_end, T)


class DiffusionSchedule:
    def __init__(self, T: int = 300, device="cpu"):
        self.T = T
        self.betas = linear_beta_schedule(T).to(device)
        self.alphas = 1.0 - self.betas
        self.alpha_bars = torch.cumprod(self.alphas, dim=0)

    def q_sample(self, x0, t, noise):
        # x0: (B, C, H, W), t: (B,) integer timesteps
        ab = self.alpha_bars[t]
        sqrt_ab = ab.sqrt().view(-1, 1, 1, 1)
        sqrt_omab = (1 - ab).sqrt().view(-1, 1, 1, 1)
        return sqrt_ab * x0 + sqrt_omab * noise


# ---- 2. Sinusoidal timestep embedding (same idea as Ch. 21) ----

class TimeEmbedding(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        self.dim = dim
        self.mlp = nn.Sequential(
            nn.Linear(dim, dim), nn.SiLU(), nn.Linear(dim, dim)
        )

    def forward(self, t: torch.Tensor):
        half = self.dim // 2
        i = torch.arange(half, device=t.device)
        freqs = torch.exp(-math.log(10000) * i / half)
        args = t.float().unsqueeze(1) * freqs.unsqueeze(0)
        emb = torch.cat([torch.sin(args), torch.cos(args)], dim=-1)
        return self.mlp(emb)


# ---- 3. A small convolutional denoiser (simplified U-Net) ----

class SimpleDenoiser(nn.Module):
    def __init__(self, time_dim: int = 64):
        super().__init__()
        self.time_embed = TimeEmbedding(time_dim)

        self.down1 = nn.Conv2d(1, 32, 3, padding=1)
        self.down2 = nn.Conv2d(32, 64, 3, stride=2, padding=1)  # ->14
        self.down3 = nn.Conv2d(64, 128, 3, stride=2, padding=1)  # ->7

        self.time_proj = nn.Linear(time_dim, 128)

        self.up1 = nn.ConvTranspose2d(
            128, 64, 4, stride=2, padding=1
        )  # 7 -> 14
        self.up2 = nn.ConvTranspose2d(
            64, 32, 4, stride=2, padding=1
        )  # 14 -> 28
        self.out = nn.Conv2d(32, 1, 3, padding=1)

    def forward(self, x, t):
        temb = self.time_embed(t)  # (B, time_dim)

        h1 = F.silu(self.down1(x))   # (B, 32, 28, 28)
        h2 = F.silu(self.down2(h1))  # (B, 64, 14, 14)
        h3 = F.silu(self.down3(h2))  # (B, 128, 7, 7)

        h3 = h3 + self.time_proj(temb).view(-1, 128, 1, 1)

        u1 = F.silu(self.up1(h3)) + h2  # skip connection
        u2 = F.silu(self.up2(u1)) + h1  # skip connection
        return self.out(u2)  # predicted noise, same shape as x


# ---- 4. Training loop ----

def train(epochs=5, T=300, batch_size=128, device="cpu"):
    tfm = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,)),
    ])
    train_set = datasets.MNIST(
        root="./data", train=True, download=True, transform=tfm
    )
    loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)

    schedule = DiffusionSchedule(T=T, device=device)
    model = SimpleDenoiser().to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=2e-4)

    for epoch in range(epochs):
        for x0, _ in loader:
            x0 = x0.to(device)
            b = x0.size(0)
            t = torch.randint(0, T, (b,), device=device)
            noise = torch.randn_like(x0)
            x_t = schedule.q_sample(x0, t, noise)

            pred_noise = model(x_t, t)
            loss = F.mse_loss(pred_noise, noise)  # eq. 30.2

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        print(f"epoch {epoch}  loss {loss.item():.4f}")

    return model, schedule


# ---- 5. Sampling ----

@torch.no_grad()
def sample(model, schedule, n_samples=16, device="cpu"):
    model.eval()
    x = torch.randn(n_samples, 1, 28, 28, device=device)
    for t_ in reversed(range(schedule.T)):
        t = torch.full(
            (n_samples,), t_, device=device, dtype=torch.long
        )
        alpha_t = schedule.alphas[t_]
        alpha_bar_t = schedule.alpha_bars[t_]
        beta_t = schedule.betas[t_]

        pred_noise = model(x, t)
        coef = beta_t / (1 - alpha_bar_t).sqrt()
        mean = (x - coef * pred_noise) / alpha_t.sqrt()

        if t_ > 0:
            z = torch.randn_like(x)
            x = mean + beta_t.sqrt() * z
        else:
            x = mean
    model.train()
    return x


if __name__ == "__main__":
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model, schedule = train(epochs=5, T=300, device=device)
    samples = sample(model, schedule, n_samples=16, device=device)
    # samples: (16, 1, 28, 28), values roughly in [-1, 1];
    # rescale to [0, 1] before viewing or saving.
