"""Portable helpers from the original conditional MNIST DDPM notebook.

The blog notebook inlines these definitions; importing this module requires
PyTorch, torchvision, and matplotlib, but no Google Drive mount.
"""

import math

import torch
import torch.nn as nn
import torch.nn.functional as F

from torchvision.utils import make_grid

import matplotlib.pyplot as plt

def show_images(images, labels=None, nrow=8, title=None):
    """
    images: tensor [B, C, H, W], value range [-1, 1]
    """
    images = images.detach().cpu()
    images = (images + 1.0) / 2.0
    images = images.clamp(0, 1)

    grid = make_grid(images, nrow=nrow)

    plt.figure(figsize=(8, 8))
    plt.imshow(grid.permute(1, 2, 0).squeeze(), cmap="gray")

    if title is not None:
        plt.title(title)

    plt.axis("off")
    plt.show()

def extract(a, t, x_shape):
    """
    Extract values from a 1-D tensor `a` at indices `t`,
    then reshape for broadcasting with image tensor.

    a: shape [T]
    t: shape [B]
    x_shape: usually [B, C, H, W]

    return: shape [B, 1, 1, 1]
    """
    B = t.shape[0]
    out = a.gather(0, t)
    return out.reshape(B, *((1,) * (len(x_shape) - 1)))


def make_noise_schedule(T, beta_start=1e-4, beta_end=0.02, device="cpu"):
    """
    Create a linear noise schedule.

    Returns a dictionary containing:
    betas:                  [T]
    alphas:                 [T]
    alpha_bars:             [T]
    sqrt_alpha_bars:        [T]
    sqrt_one_minus_alpha_bars: [T]
    sqrt_recip_alphas:      [T]
    """
    betas = torch.linspace(beta_start, beta_end, T, device=device)

    alphas = 1.0 - betas

    alpha_bars = torch.cumprod(alphas, dim=0)

    sqrt_alpha_bars = torch.sqrt(alpha_bars)

    sqrt_one_minus_alpha_bars = torch.sqrt(1.0 - alpha_bars)

    sqrt_recip_alphas = torch.sqrt(1.0 / alphas)

    return {
        "betas": betas,
        "alphas": alphas,
        "alpha_bars": alpha_bars,
        "sqrt_alpha_bars": sqrt_alpha_bars,
        "sqrt_one_minus_alpha_bars": sqrt_one_minus_alpha_bars,
        "sqrt_recip_alphas": sqrt_recip_alphas,
    }


def add_noise(x0, t, noise, schedule):
    """
    Add noise to clean images x0 at timestep t.

    x0:    [B, C, H, W]
    t:     [B]
    noise: [B, C, H, W]

    return:
    x_t: noisy image at timestep t
    """
    sqrt_alpha_bar_t = extract(
        schedule["sqrt_alpha_bars"],
        t,
        x0.shape,
    )

    sqrt_one_minus_alpha_bar_t = extract(
        schedule["sqrt_one_minus_alpha_bars"],
        t,
        x0.shape,
    )

    x_t = sqrt_alpha_bar_t * x0 + sqrt_one_minus_alpha_bar_t * noise

    return x_t

def plot_noise_schedule(schedule, T):
    import matplotlib.pyplot as plt
    import torch
    t = torch.arange(T).cpu().numpy()

    plt.figure(figsize=(12, 5))

    plt.subplot(1, 3, 1)
    plt.plot(t, schedule['betas'].cpu().numpy())
    plt.title('Betas Schedule')
    plt.xlabel('Time (t)')
    plt.ylabel('Beta Value')
    plt.grid(True)

    plt.subplot(1, 3, 2)
    plt.plot(t, schedule['alphas'].cpu().numpy())
    plt.title('Alphas Schedule')
    plt.xlabel('Time (t)')
    plt.ylabel('Alpha Value')
    plt.grid(True)

    plt.subplot(1, 3, 3)
    plt.plot(t, schedule['alpha_bars'].cpu().numpy())
    plt.title('Alpha Bars Schedule')
    plt.xlabel('Time (t)')
    plt.ylabel('Alpha Bar Value')
    plt.grid(True)

    plt.tight_layout()
    plt.show()

def visualize_forward_diffusion(x0, schedule, n_steps=10):
    """Show independent forward marginals with a fixed display range."""
    import torch
    import matplotlib.pyplot as plt
    # Pick a few time steps to show
    indices = torch.linspace(0, len(schedule['betas']) - 1, n_steps).long()

    plt.figure(figsize=(15, 3))
    for i, t in enumerate(indices):
        # Forward diffusion formula: q(xt|x0)
        sqrt_alpha_bar = schedule['sqrt_alpha_bars'][t]
        sqrt_one_minus_alpha_bar = schedule['sqrt_one_minus_alpha_bars'][t]

        noise = torch.randn_like(x0[0:1])
        xt = sqrt_alpha_bar * x0[0:1] + sqrt_one_minus_alpha_bar * noise

        plt.subplot(1, n_steps, i + 1)
        img = xt.squeeze().cpu().numpy()
        plt.imshow((img + 1.0) / 2.0, cmap="gray", vmin=0, vmax=1)
        plt.title(f't={t.item()}')
        plt.axis('off')
    plt.show()


def sinusoidal_time_embedding(t, dim):
    """
    Create sinusoidal timestep embeddings.

    t:   [B], integer timesteps
    dim: embedding dimension

    return:
    emb: [B, dim]
    """
    if dim < 4:
        raise ValueError("Use an embedding dimension of at least 4.")
    half_dim = dim // 2

    device = t.device

    # Frequencies from high to low
    exponent = -math.log(10000) * torch.arange(
        half_dim,
        device=device,
        dtype=torch.float32,
    ) / (half_dim - 1)

    freqs = torch.exp(exponent)

    # [B, 1] * [1, half_dim] -> [B, half_dim]
    args = t.float()[:, None] * freqs[None, :]

    emb = torch.cat([torch.sin(args), torch.cos(args)], dim=-1)

    # If dim is odd, pad one dimension
    if dim % 2 == 1:
        emb = F.pad(emb, (0, 1))

    return emb

class ConditionEmbedding(nn.Module):
    def __init__(self, num_classes, time_emb_dim):
        super().__init__()

        self.time_mlp = nn.Sequential(
            nn.Linear(time_emb_dim, time_emb_dim * 4),
            nn.SiLU(),
            nn.Linear(time_emb_dim * 4, time_emb_dim),
        )

        self.label_emb = nn.Embedding(num_classes, time_emb_dim)

    def forward(self, t, y):
        """
        t: [B]
        y: [B]

        return:
        emb: [B, time_emb_dim]
        """
            
        t_emb = sinusoidal_time_embedding(t, self.label_emb.embedding_dim)
        t_emb = self.time_mlp(t_emb)

        y_emb = self.label_emb(y)

        emb = t_emb + y_emb

        return emb
    
class ResBlock(nn.Module):
    def __init__(self, in_ch, out_ch, emb_dim):
        super().__init__()

        self.conv1 = nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1)

        self.emb_proj = nn.Linear(emb_dim, out_ch)

        self.act = nn.SiLU()

        if in_ch != out_ch:
            self.skip = nn.Conv2d(in_ch, out_ch, kernel_size=1)
        else:
            self.skip = nn.Identity()

    def forward(self, x, emb):
        """
        x:   [B, in_ch, H, W]
        emb: [B, emb_dim]

        return:
        out: [B, out_ch, H, W]
        """
        h = self.conv1(x)

        emb_out = self.emb_proj(emb)
        emb_out = emb_out[:, :, None, None]

        h = h + emb_out
        h = self.act(h)

        h = self.conv2(h)
        h = self.act(h)

        return h + self.skip(x)

class Downsample(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv = nn.Conv2d(
            channels,
            channels,
            kernel_size=4,
            stride=2,
            padding=1,
        )

    def forward(self, x):
        return self.conv(x)

class Upsample(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv = nn.Conv2d(
            channels,
            channels,
            kernel_size=3,
            padding=1,
        )

    def forward(self, x):
        x = F.interpolate(x, scale_factor=2, mode="nearest")
        x = self.conv(x)
        return x

class ConditionalUNet(nn.Module):
    def __init__(
        self,
        image_channels=1,
        base_channels=64,
        num_classes=10,
        time_emb_dim=128,
    ):
        super().__init__()

        self.cond_emb = ConditionEmbedding(
            num_classes=num_classes,
            time_emb_dim=time_emb_dim,
        )

        # Encoder
        self.down_block1 = ResBlock(
            in_ch=image_channels,
            out_ch=base_channels,
            emb_dim=time_emb_dim,
        )
        self.downsample1 = Downsample(base_channels)

        self.down_block2 = ResBlock(
            in_ch=base_channels,
            out_ch=base_channels * 2,
            emb_dim=time_emb_dim,
        )
        self.downsample2 = Downsample(base_channels * 2)

        # Bottleneck
        self.mid_block = ResBlock(
            in_ch=base_channels * 2,
            out_ch=base_channels * 2,
            emb_dim=time_emb_dim,
        )

        # Decoder
        self.upsample1 = Upsample(base_channels * 2)

        self.up_block1 = ResBlock(
            in_ch=base_channels * 4,   # concat: 128 + 128
            out_ch=base_channels,
            emb_dim=time_emb_dim,
        )

        self.upsample2 = Upsample(base_channels)

        self.up_block2 = ResBlock(
            in_ch=base_channels * 2,   # concat: 64 + 64
            out_ch=base_channels,
            emb_dim=time_emb_dim,
        )

        self.out_conv = nn.Conv2d(
            base_channels,
            image_channels,
            kernel_size=1,
        )

    def forward(self, x, t, y):
        """
        x: [B, 1, 28, 28]
        t: [B]
        y: [B]

        return:
        predicted noise: [B, 1, 28, 28]
        """
        emb = self.cond_emb(t, y)

        # Encoder
        h1 = self.down_block1(x, emb)      # [B, 64, 28, 28]
        h = self.downsample1(h1)           # [B, 64, 14, 14]

        h2 = self.down_block2(h, emb)      # [B, 128, 14, 14]
        h = self.downsample2(h2)           # [B, 128, 7, 7]

        # Bottleneck
        h = self.mid_block(h, emb)         # [B, 128, 7, 7]

        # Decoder
        h = self.upsample1(h)              # [B, 128, 14, 14]
        h = torch.cat([h, h2], dim=1)      # [B, 256, 14, 14]
        h = self.up_block1(h, emb)         # [B, 64, 14, 14]

        h = self.upsample2(h)              # [B, 64, 28, 28]
        h = torch.cat([h, h1], dim=1)      # [B, 128, 28, 28]
        h = self.up_block2(h, emb)         # [B, 64, 28, 28]

        out = self.out_conv(h)             # [B, 1, 28, 28]

        return out


def diffusion_loss(model, x0, y, schedule, cfg):
    """
    Compute DDPM noise prediction loss.

    x0: [B, 1, 28, 28], clean MNIST images, range [-1, 1]
    y:  [B], digit labels

    return:
    loss: scalar
    """
    B = x0.shape[0]

    # 1. Randomly sample timestep t for each image in the batch
    t = torch.randint(
        low=0,
        high=cfg.T,
        size=(B,),
        device=x0.device,
    )

    # 2. Sample Gaussian noise
    noise = torch.randn_like(x0)

    # 3. Create noisy image x_t
    x_t = add_noise(x0, t, noise, schedule)

    # 4. Predict the noise
    pred_noise = model(x_t, t, y)

    # 5. MSE loss between predicted noise and real noise
    loss = F.mse_loss(pred_noise, noise)

    return loss

def train(model, train_loader, schedule, cfg, device):
    optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.lr)

    model.train()

    for epoch in range(cfg.epochs):
        total_loss = 0.0

        for step, (x0, y) in enumerate(train_loader):
            x0 = x0.to(device)
            y = y.to(device)

            loss = diffusion_loss(model, x0, y, schedule, cfg)

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

            total_loss += loss.item()

            if step % 100 == 0:
                print(
                    f"Epoch [{epoch+1}/{cfg.epochs}] "
                    f"Step [{step}/{len(train_loader)}] "
                    f"Loss: {loss.item():.4f}"
                )

        avg_loss = total_loss / len(train_loader)

        print(
            f"Epoch [{epoch+1}/{cfg.epochs}] "
            f"Average Loss: {avg_loss:.4f}"
        )

    return model

@torch.no_grad()
def p_sample(model, x, t, y, schedule):
    """
    Sample x_{t-1} from x_t.

    model: conditional noise predictor
    x:     [B, 1, 28, 28], current noisy image x_t
    t:     [B], current timestep
    y:     [B], digit labels

    return:
    x_prev: [B, 1, 28, 28]
    """
    betas_t = extract(schedule["betas"], t, x.shape)
    sqrt_one_minus_alpha_bars_t = extract(
        schedule["sqrt_one_minus_alpha_bars"],
        t,
        x.shape,
    )
    sqrt_recip_alphas_t = extract(
        schedule["sqrt_recip_alphas"],
        t,
        x.shape,
    )

    # Predict noise epsilon_theta(x_t, t, y)
    pred_noise = model(x, t, y)

    # DDPM reverse mean
    model_mean = sqrt_recip_alphas_t * (
        x - betas_t * pred_noise / sqrt_one_minus_alpha_bars_t
    )

    # If t == 0, return the mean directly
    # If t > 0, add Gaussian noise
    noise = torch.randn_like(x)

    nonzero_mask = (t != 0).float().reshape(
        x.shape[0],
        *((1,) * (len(x.shape) - 1)),
    )

    x_prev = model_mean + nonzero_mask * torch.sqrt(betas_t) * noise

    return x_prev

@torch.no_grad()
def sample_digits(model, digits, schedule, cfg, device):
    """
    Generate one image for each digit in `digits`.

    digits: list of int, e.g. [0, 1, 2, ..., 9]

    return:
    x: [B, 1, 28, 28], generated images
    """
    model.eval()

    B = len(digits)

    # Start from pure Gaussian noise
    x = torch.randn(
        B,
        cfg.channels,
        cfg.image_size,
        cfg.image_size,
        device=device,
    )

    y = torch.tensor(digits, device=device, dtype=torch.long)

    # Reverse diffusion: T-1 -> 0
    for i in reversed(range(cfg.T)):
        t = torch.full(
            (B,),
            i,
            device=device,
            dtype=torch.long,
        )

        x = p_sample(model, x, t, y, schedule)

    return x