# 使得在 notebook 中顯示繪圖，而不是在外部窗口中顯示
%matplotlib inline  
import matplotlib.pyplot as plt
import numpy as np
from sklearn.datasets import make_s_curve
import torch
import torch.nn as nn
import io
from PIL import Image

s_curve, _ = make_s_curve(10 ** 4, noise=0.1)

s_curve = s_curve[:, [0, 2]] / 10.
print(F"shape of Moons:{np.shape(s_curve)}")

data = s_curve.T
 
# 繪制 S 形曲線數(shù)據(jù)集
fig, ax = plt.subplots()
ax.scatter(*data, color='red', edgecolor='white')
ax.axis('off')

dataset = torch.Tensor(s_curve).float()

num_step = 100  # 一開始可以由beta、分布的均值和標準差來共同確定
 
# 指定每一步的beta
betas = torch.linspace(-6, 6, num_step)
betas = torch.sigmoid(betas) * (0.5e-2 - 1e-5) + 1e-5
 
# 計算alpha、alpha_prod、alpha_prod_previous、alpha_bar_sqrt等變量的值
alphas = 1 - betas
alphas_prod = torch.cumprod(alphas, dim=0)
alphas_prod_p = torch.cat([torch.tensor([1]).float(), alphas_prod[:-1]], 0)  # p表示previous
alphas_bar_sqrt = torch.sqrt(alphas_prod)
one_minus_alphas_bar_log = torch.log(1 - alphas_prod)
one_minus_alphas_bar_sqrt = torch.sqrt(1 - alphas_prod)
 
assert alphas.shape == alphas_prod.shape == alphas_prod_p.shape == alphas_bar_sqrt.shape == one_minus_alphas_bar_log.shape == one_minus_alphas_bar_sqrt.shape
print(f"all the same shape:{betas.shape}")

# 計算任意時刻的x的采樣值，基于x_0和參數(shù)重整化技巧
def q_x(x_0, t):
    """可以基于x[0]得到任意時刻t的x[t]"""
    noise = torch.randn_like(x_0)  # noise是從正態(tài)分布中生成的隨機噪聲
    alphas_t = alphas_bar_sqrt[t]
    alphas_l_m_t = one_minus_alphas_bar_sqrt[t]
    return (alphas_t * x_0 + alphas_l_m_t * noise) # 在x[0]的基礎(chǔ)上添加噪聲

num_shows = 20
fig, axs = plt.subplots(2, 10, figsize=(28, 7))
plt.rc('text', color='blue')
# 共有10000個點，每個點包含兩個坐標
# 生成100步以內(nèi)每隔5步加噪聲后的圖像
for i in range(num_shows):
    j = i // 10
    k = i % 10
    q_i = q_x(dataset, torch.tensor([i * num_step // num_shows]))  # 生成t時刻的采樣數(shù)據(jù)
    axs[j, k].scatter(q_i[:, 0], q_i[:, 1], color='red', edgecolor='white')
    
    axs[j, k].set_axis_off()
    axs[j, k].set_title('$q(\mathbf{x}_{'+ str(i * num_step // num_shows) + '})$')

class MLPDiffusion(nn.Module):
    def __init__(self, n_steps, num_units=128):
        super(MLPDiffusion, self).__init__()
        
        self.linears = nn.ModuleList(
        [
            nn.Linear(2, num_units),
            nn.ReLU(),
            nn.Linear(num_units, num_units),
            nn.ReLU(),
            nn.Linear(num_units, num_units),
            nn.ReLU(),
            nn.Linear(num_units, 2),
        ])
        
        self.step_embeddings = nn.ModuleList(
        [
            nn.Embedding(n_steps, num_units),
            nn.Embedding(n_steps, num_units),
            nn.Embedding(n_steps, num_units),
        ])
    
    def forward(self, x, t):
        for idx, embedding_layer in enumerate(self.step_embeddings):
            t_embedding = embedding_layer(t)
            x = self.linears[2 * idx](x)
            x += t_embedding
            x = self.linears[2 * idx + 1](x)
        
        x = self.linears[-1](x)
        return x

def diffusion_loss_fn(model, x_0, alphas_bar_sqrt, one_minus_alphas_bar_sqrt, n_steps):
    """
    對任意時刻t進行采樣計算loss
    param：
    model：模型 
    x_0：初始狀態(tài) 
    alphas_bar_sqrt、one_minus_alphas_bar_sqrt： 參數(shù) 
    n_steps：時間步數(shù)
    return：損失值
    """
    batch_size = x_0.shape[0]
    # 隨機采樣一個時刻t，為了提高訓練效率，這里確保t不重復(fù)
    # 對一個batchsize樣本生成隨機的時刻t，覆蓋到更多不同的t
    t = torch.randint(0, n_steps, size=(batch_size // 2,))
    t = torch.cat([t, n_steps - 1 - t], dim=0)  # [batch]
    t = t.unsqueeze(-1)  # [batch, 1]
    # x0的系數(shù)
    a = alphas_bar_sqrt[t]
    # eps的系數(shù)
    aml = one_minus_alphas_bar_sqrt[t]
    # 生成隨機噪聲eps
    e = torch.randn_like(x_0)
    # 構(gòu)造模型的輸入
    x = x_0 * a + e * aml
    # 送入模型，得到t時刻的隨機噪聲預(yù)測值
    output = model(x, t.squeeze(-1))
    # 與真實噪聲一起計算誤差，求平均值
    return (e - output).square().mean()

def p_sample_loop(model, shape, n_step, betas, one_minus_alphas_bar_sqrt):
    """從x[T]恢復(fù)x[T - 1]、x[T - 2]、...、x[0]"""
    cur_x = torch.randn(shape)
    x_seq = [cur_x]
    for i in reversed(range(n_step)):
        cur_x = p_sample(model, cur_x, i, betas, one_minus_alphas_bar_sqrt)
        x_seq.append(cur_x)
    return x_seq
 
def p_sample(model, x, t, betas, one_minus_alphas_bar_sqrt):
    """從x[T]采樣t時刻的重構(gòu)值"""
    t = torch.tensor([t])
    coeff = betas[t] / one_minus_alphas_bar_sqrt[t]
    eps_theta = model(x, t)
    mean = (1 / (1 - betas[t]).sqrt()) * (x - (coeff * eps_theta))
    z = torch.randn_like(x)
    sigma_t = betas[t].sqrt()
    sample = mean + sigma_t * z
    return (sample)

seed = 1234  # 確保程序在每次運行時生成的隨機數(shù)序列都是一樣的
 
class EMA():
    """構(gòu)建一個參數(shù)平滑器，以便更好地泛化模型并減少過擬合"""
    def __init__(self, mu=0.01):
        self.mu = mu
        self.shadow = {}
        
    def register(self, name, val):
        self.shadow[name] = val.clone()
        
    def __call__(self, name, x):
        assert name in self.shadow
        new_average = self.mu * x + (1.0 - self.mu) * self.shadow[name]
        return new_average
print('Training model.....')
 
 
batch_size = 512  # 批訓練大小
dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)
num_epoch = 4000  # 定義迭代4000次
plt.rc('text', color='blue')
 
model = MLPDiffusion(num_step)  # 輸出維度是2，輸入是x和step
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
 
for t in range(num_epoch):
    for idx, batch_x in enumerate(dataloader):
        loss = diffusion_loss_fn(model, batch_x, alphas_bar_sqrt, one_minus_alphas_bar_sqrt, num_step)
        optimizer.zero_grad()  # 對梯度進行清零，防止網(wǎng)絡(luò)權(quán)重更新過于迅速或不穩(wěn)定，無法得到正確的收斂結(jié)果
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), 1.)  # 對梯度進行裁剪，避免出現(xiàn)梯度爆炸
        optimizer.step()
    if (t % 100 == 0):
        print(loss)
        x_seq = p_sample_loop(model, dataset.shape, num_step, betas, one_minus_alphas_bar_sqrt)  # 共有100個元素
        
        fig, axs = plt.subplots(1, 5, figsize=(28, 7))
        for i in range(1, 6):
            cur_x = x_seq[i * 20].detach()
            axs[i - 1].scatter(cur_x[:, 0], cur_x[:, 1], color='red', edgecolor='white')
            axs[i - 1].set_axis_off()
            axs[i - 1].set_title('$q(\mathbf{x}_{'+str(i * 20)+'})$')

# 生成前向過程，也就是逐步加噪聲
imgs = []
for i in range(100):
    plt.clf()
    torch_i = q_x(dataset, torch.tensor([i]))
    plt.scatter(q_i[:, 0], q_i[:, 1], color='red', edgecolor='white', s=5)
    plt.axis('off')
    
    img_buf = io.BytesIO()
    plt.savefig(img_buf, format='png')
    img = Image.open(img_buf)
    imgs.append(img)

# 生成逆過程，也就是逐步復(fù)原
reverse = []
for i in range(100):
    plt.clf()
    cur_x = x_seq[i].detach()  # 拿到訓練末尾階段生成的x_seq
    plt.scatter(cur_x[:, 0], cur_x[:, 1], color='red', edgecolor='white', s=5)
    plt.axis('off')
    
    img_buf = io.BytesIO()
    plt.savefig(img_buf, format='png')
    img = Image.open(img_buf)
    reverse.append(img)

imgs = imgs + reverse
 
imgs[0].save("diffusion.gif", format='GIF', append_images=imgs, save_all=True, duration=100, loop=0)