import os
import argparse
import math
import torch
from torch.nn import functional as F
from torchvision import transforms
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from PIL import Image

from main import UNet


def build_schedule(timesteps, device):
    betas = torch.linspace(1e-4, 0.02, timesteps).to(device)
    alphas = 1.0 - betas
    alpha_bar = torch.cumprod(alphas, dim=0)
    return betas, alphas, alpha_bar


def denoise_step(model, x, t, alphas, alpha_bar, betas):
    n = x.size(0)
    t_batch = torch.full((n,), t, device=x.device, dtype=torch.long)
    eps_pred = model(x, t_batch)
    alpha_t = alphas[t]
    alpha_bar_t = alpha_bar[t]
    x = (x - (1 - alpha_t) / torch.sqrt(1 - alpha_bar_t) * eps_pred) / torch.sqrt(alpha_t)
    if t > 0:
        x += torch.sqrt(betas[t]) * torch.randn_like(x)
    return x


def to_grid_image(x, ncols, cmap="gray"):
    n = x.size(0)
    nrows = math.ceil(n / ncols)
    fig, axes = plt.subplots(nrows, ncols, figsize=(ncols * 2, nrows * 2))
    axes = np.atleast_1d(axes).flatten()
    for i, ax in enumerate(axes):
        if i < n:
            ax.imshow(x[i, 0], cmap=cmap)
        ax.axis("off")
    plt.subplots_adjust(wspace=0.05, hspace=0.05)
    fig.canvas.draw()
    buf = np.asarray(fig.canvas.buffer_rgba())[:, :, :3]
    plt.close(fig)
    return Image.fromarray(buf)


def main():
    parser = argparse.ArgumentParser(description="DDPM 采样过程逐帧记录")
    parser.add_argument("--ckpt", type=str, default="model_state_dict.pth")
    parser.add_argument("--timesteps", type=int, default=600)
    parser.add_argument("--samples", type=int, default=4)
    parser.add_argument("--outdir", type=str, default="sampling_frames")
    parser.add_argument("--gif", action="store_true", help="额外输出 GIF")
    parser.add_argument("--gif-fps", type=int, default=20)
    args = parser.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"

    model = UNet().to(device)
    state = torch.load(args.ckpt, map_location=device)
    if isinstance(state, dict) and "state_dict" in state:
        state = state["state_dict"]
    model.load_state_dict(state)
    model.eval()

    betas, alphas, alpha_bar = build_schedule(args.timesteps, device)

    os.makedirs(args.outdir, exist_ok=True)

    x = torch.randn(args.samples, 1, 28, 28, device=device)

    frames = []
    with torch.no_grad():
        for t in reversed(range(args.timesteps)):
            x = denoise_step(model, x, t, alphas, alpha_bar, betas)
            grid = x.cpu().clamp(-1, 1)
            grid = (grid + 1) / 2
            img = to_grid_image(grid, ncols=2)
            path = os.path.join(args.outdir, f"step_{t:04d}.png")
            img.save(path)
            frames.append(path)
            if t % 50 == 0 or t == 0:
                print(f"step {t} 已保存 -> {path}")

    print(f"全部 {args.timesteps} 帧已保存到 {args.outdir}/")

    if args.gif:
        images = [Image.open(p) for p in frames]
        gif_path = os.path.join(args.outdir, "sampling.gif")
        images[0].save(
            gif_path,
            save_all=True,
            append_images=images[1:],
            duration=1000 // args.gif_fps,
            loop=0,
        )
        print(f"GIF 已保存 -> {gif_path}")


if __name__ == "__main__":
    main()