"""Gaussian diffusion utilities for Phase 4 (DDPM).

Provides noise schedules, the forward (noising) process, training loss,
and DDIM deterministic sampling (Song et al., 2020).

Convention: alpha_bars is a 1-D tensor of length T, where alpha_bars[t]
= ᾱ_{t+1} in 1-indexed notation.  Timestep t used in the training loop
is a 0-indexed integer in [0, T).  At t=0 the image is almost clean
(ᾱ ≈ 1 − β_1); at t=T−1 the image is almost pure noise (ᾱ ≈ 0).
"""
import math

import torch
import torch.nn.functional as F


# ── Noise schedules ──────────────────────────────────────────────────────────

def linear_betas(T: int, beta_start: float = 1e-4, beta_end: float = 0.02) -> torch.Tensor:
    """Ho et al. (2020) linear schedule."""
    return torch.linspace(beta_start, beta_end, T)


def cosine_betas(T: int, s: float = 0.008) -> torch.Tensor:
    """Nichol & Dhariwal (2021) cosine schedule — avoids over-denoising at low t."""
    t = torch.linspace(0, T, T + 1)
    f = torch.cos((t / T + s) / (1 + s) * math.pi / 2) ** 2
    alpha_bar = f / f[0]
    betas = 1 - alpha_bar[1:] / alpha_bar[:-1]
    return betas.clamp(max=0.999)


def make_alpha_bars(betas: torch.Tensor) -> torch.Tensor:
    """Cumulative product of (1 − β), shape (T,)."""
    return (1.0 - betas).cumprod(0)


# ── Forward process ──────────────────────────────────────────────────────────

def q_sample(
    x0: torch.Tensor,
    t: torch.Tensor,
    alpha_bars: torch.Tensor,
    noise: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Add noise to x0 at timestep t.  Returns (x_t, noise)."""
    if noise is None:
        noise = torch.randn_like(x0)
    ab   = alpha_bars[t].to(x0.device)[:, None, None, None]
    x_t  = ab.sqrt() * x0 + (1 - ab).sqrt() * noise
    return x_t, noise


# ── Training loss ────────────────────────────────────────────────────────────

def diffusion_loss(
    model,
    x0: torch.Tensor,
    t: torch.Tensor,
    alpha_bars: torch.Tensor,
    pred_type: str = "eps",
) -> torch.Tensor:
    """MSE on the model's prediction vs the true target.

    pred_type="eps" → target is the added noise ε (Ho et al.)
    pred_type="v"   → target is v = √ᾱ·ε − √(1−ᾱ)·x0  (Salimans & Ho)
    """
    x_t, noise = q_sample(x0, t, alpha_bars)
    pred = model(x_t, t)

    if pred_type == "eps":
        target = noise
    else:  # v
        ab = alpha_bars[t].to(x0.device)[:, None, None, None]
        target = ab.sqrt() * noise - (1 - ab).sqrt() * x0

    return F.mse_loss(pred, target)


# ── DDIM deterministic sampling ───────────────────────────────────────────────

@torch.no_grad()
def ddim_sample(
    model,
    n: int,
    image_size: int,
    alpha_bars: torch.Tensor,
    n_steps: int = 100,
    pred_type: str = "eps",
    device: str = "cuda",
    batch_size: int = 32,
) -> torch.Tensor:
    """Generate n images via DDIM (eta=0, deterministic).

    Batches internally to avoid OOM when n is large.
    Returns tensor shape (n, 3, image_size, image_size) in [-1, 1].
    """
    model.eval()
    T = len(alpha_bars)

    # Build reversed subsequence: [T-1, T-1-step, ..., 0]
    step = max(T // n_steps, 1)
    ts   = list(range(T - 1, -1, -step))[:n_steps]
    if ts[-1] != 0:
        ts.append(0)

    results = []
    remaining = n
    while remaining > 0:
        bsz = min(batch_size, remaining)
        x   = torch.randn(bsz, 3, image_size, image_size, device=device)

        for i, t_cur in enumerate(ts):
            t_prev = ts[i + 1] if i + 1 < len(ts) else -1

            t_batch = torch.full((bsz,), t_cur, device=device, dtype=torch.long)
            ab_t    = alpha_bars[t_cur].to(device)
            ab_prev = alpha_bars[t_prev].to(device) if t_prev >= 0 else torch.ones(1, device=device)

            pred = model(x, t_batch)

            # Reconstruct x0 from prediction
            if pred_type == "eps":
                x0_hat = (x - (1 - ab_t).sqrt() * pred) / ab_t.sqrt()
            else:  # v
                x0_hat = ab_t.sqrt() * x - (1 - ab_t).sqrt() * pred
            x0_hat = x0_hat.clamp(-1, 1)

            # DDIM step
            eps_hat = (x - ab_t.sqrt() * x0_hat) / (1 - ab_t).sqrt()
            x = ab_prev.sqrt() * x0_hat + (1 - ab_prev).sqrt() * eps_hat

        results.append(x.cpu())
        remaining -= bsz

    return torch.cat(results)[:n]
