''' 
This script does conditional image generation on MNIST, using a diffusion model

This code is modified from,
https://github.com/cloneofsimo/minDiffusion

Diffusion model is based on DDPM,
https://arxiv.org/abs/2006.11239

The conditioning idea is taken from 'Classifier-Free Diffusion Guidance',
https://arxiv.org/abs/2207.12598

This technique also features in ImageGen 'Photorealistic Text-to-Image Diffusion Modelswith Deep Language Understanding',
https://arxiv.org/abs/2205.11487

'''

from tqdm import tqdm # 進捗バーを表示するためのライブラリ
import torch 
import torch.nn as nn # ニューラルネットワークのモジュールを提供
from torch.utils.data import DataLoader # データローダーを提供
from torchvision import transforms # データの変換を提供
from torchvision.datasets import MNIST # MNISTデータセットを提供
from torchvision.utils import save_image, make_grid # 画像の保存とグリッド化を提供
import matplotlib.pyplot as plt # プロット作成用のライブラリ
from matplotlib.animation import FuncAnimation, PillowWriter # アニメーション作成用のライブラリ
import numpy as np # 数値計算用のライブラリ

class ResidualConvBlock(nn.Module):
    def __init__(
        self, in_channels: int, out_channels: int, is_res: bool = False) -> None:
        super().__init__()
        '''standard ResNet style convolutional block
        is_res: (bool) 残差接続を行うかどうかの真偽値
        '''
        self.same_channels = in_channels==out_channels
        self.is_res = is_res
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 3, 1, 1),
            nn.BatchNorm2d(out_channels),
            nn.GELU(),
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(out_channels, out_channels, 3, 1, 1),
            nn.BatchNorm2d(out_channels),
            nn.GELU(),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # ニューラルネットワークの順伝播を行う
        if self.is_res:
            x1 = self.conv1(x)
            x2 = self.conv2(x1)
            # this adds on correct residual in case channels have increased
            if self.same_channels:
                out = x + x2 # 入力と2層目の出力を足す
            else:
                out = x1 + x2  # 1層目と2層目の出力を足す
            return out / 1.414
        else:
            x1 = self.conv1(x)
            x2 = self.conv2(x1)
            return x2


class UnetDown(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(UnetDown, self).__init__()
        '''
        process and downscale the image feature maps
        '''
        layers = [ResidualConvBlock(in_channels, out_channels), nn.MaxPool2d(2)]
        self.model = nn.Sequential(*layers)

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


class UnetUp(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(UnetUp, self).__init__()
        '''
        process and upscale the image feature maps
        '''
        layers = [
            nn.ConvTranspose2d(in_channels, out_channels, 2, 2),
            ResidualConvBlock(out_channels, out_channels),
            ResidualConvBlock(out_channels, out_channels),
        ]
        self.model = nn.Sequential(*layers)

    def forward(self, x, skip):
        x = torch.cat((x, skip), 1)
        x = self.model(x)
        return x


class EmbedFC(nn.Module):
    def __init__(self, input_dim, emb_dim):
        super(EmbedFC, self).__init__()
        '''
        generic one layer FC NN for embedding things  
        '''
        self.input_dim = input_dim
        layers = [
            nn.Linear(input_dim, emb_dim),
            nn.GELU(),
            nn.Linear(emb_dim, emb_dim),
        ]
        self.model = nn.Sequential(*layers)

    def forward(self, x):
        x = x.view(-1, self.input_dim)
        return self.model(x)


class ContextUnet(nn.Module):
    """ 拡散モデルのニューラルネットワーク, U-Netで構成される """
    def __init__(self, in_channels, n_feat = 256, n_classes=10):
        super(ContextUnet, self).__init__()

        self.in_channels = in_channels
        self.n_feat = n_feat
        self.n_classes = n_classes

        self.init_conv = ResidualConvBlock(in_channels, n_feat, is_res=True)

        self.down1 = UnetDown(n_feat, n_feat)
        self.down2 = UnetDown(n_feat, 2 * n_feat)

        self.to_vec = nn.Sequential(nn.AvgPool2d(7), nn.GELU())

        self.timeembed1 = EmbedFC(1, 2*n_feat)
        self.timeembed2 = EmbedFC(1, 1*n_feat)
        self.contextembed1 = EmbedFC(n_classes, 2*n_feat)
        self.contextembed2 = EmbedFC(n_classes, 1*n_feat)

        self.up0 = nn.Sequential(
            # nn.ConvTranspose2d(6 * n_feat, 2 * n_feat, 7, 7), # when concat temb and cemb end up w 6*n_feat
            nn.ConvTranspose2d(2 * n_feat, 2 * n_feat, 7, 7), # otherwise just have 2*n_feat
            nn.GroupNorm(8, 2 * n_feat),
            nn.ReLU(),
        )

        self.up1 = UnetUp(4 * n_feat, n_feat)
        self.up2 = UnetUp(2 * n_feat, n_feat)
        self.out = nn.Sequential(
            nn.Conv2d(2 * n_feat, n_feat, 3, 1, 1),
            nn.GroupNorm(8, n_feat),
            nn.ReLU(),
            nn.Conv2d(n_feat, self.in_channels, 3, 1, 1),
        )

    def forward(self, x, c, t, context_mask):
        # x is (noisy) image, c is context label, t is timestep, 
        # context_mask says which samples to block the context on

        x = self.init_conv(x)
        down1 = self.down1(x)
        down2 = self.down2(down1)
        hiddenvec = self.to_vec(down2)

        # convert context to one hot embedding
        c = nn.functional.one_hot(c, num_classes=self.n_classes).type(torch.float)
        
        # mask out context if context_mask == 1
        context_mask = context_mask[:, None]
        context_mask = context_mask.repeat(1,self.n_classes)
        context_mask = (-1*(1-context_mask)) # need to flip 0 <-> 1
        c = c * context_mask
        
        # embed context, time step
        cemb1 = self.contextembed1(c).view(-1, self.n_feat * 2, 1, 1)
        temb1 = self.timeembed1(t).view(-1, self.n_feat * 2, 1, 1)
        cemb2 = self.contextembed2(c).view(-1, self.n_feat, 1, 1)
        temb2 = self.timeembed2(t).view(-1, self.n_feat, 1, 1)

        # could concatenate the context embedding here instead of adaGN
        # hiddenvec = torch.cat((hiddenvec, temb1, cemb1), 1)

        up1 = self.up0(hiddenvec)
        # up2 = self.up1(up1, down2) # if want to avoid add and multiply embeddings
        up2 = self.up1(cemb1*up1+ temb1, down2)  # add and multiply embeddings
        up3 = self.up2(cemb2*up2+ temb2, down1)
        out = self.out(torch.cat((up3, x), 1))
        return out


def ddpm_schedules(beta1, beta2, T):
    """ 拡散モデルの係数を dict で返す関数 """
    assert beta1 < beta2 < 1.0, "beta1 and beta2 must be in (0, 1)"
    
    beta_t = (beta2 - beta1) * torch.arange(0, T + 1, dtype=torch.float32) / T + beta1 # 拡散モデルのハイパーパラメータ
    sqrt_beta_t = torch.sqrt(beta_t) # 拡散過程の標準偏差
    alpha_t = 1 - beta_t # 拡散過程の平均のスケーリングの2乗
    log_alpha_t = torch.log(alpha_t)
    alphabar_t = torch.cumsum(log_alpha_t, dim=0).exp()

    sqrtmab = torch.sqrt(1 - alphabar_t) # 拡散過程の一般項の標準偏差
    sqrtab = torch.sqrt(alphabar_t) # 拡散過程の一般項の平均のスケーリング
    
    oneover_sqrta = 1 / torch.sqrt(alpha_t) # 逆拡散過程に使用する係数
    mab_over_sqrtmab_inv = (1 - alpha_t) / sqrtmab # 逆拡散過程に使用する係数

    return {
        "sqrt_beta_t": sqrt_beta_t,  
        "alpha_t": alpha_t,  
        "alphabar_t": alphabar_t,  
        "sqrtmab": sqrtmab,    
        "sqrtab": sqrtab,  
        "oneover_sqrta": oneover_sqrta,  
        "mab_over_sqrtmab": mab_over_sqrtmab_inv,  
    }


class DDPM(nn.Module):
    def __init__(self, nn_model, betas, n_T, device, drop_prob=0.1):
        """ 拡散モデルのクラス """
        super(DDPM, self).__init__()
        self.nn_model = nn_model.to(device) # ニューラルネットワーク

        for k, v in ddpm_schedules(betas[0], betas[1], n_T).items():
            self.register_buffer(k, v) # モデルのバッファに保存

        self.n_T = n_T # 総時間ステップ数
        self.device = device # PyTorch のデバイス設定
        self.drop_prob = drop_prob # 条件無し生成を学習する確率
        self.loss_mse = nn.MSELoss() # 損失関数

    def forward_process(self, x, n_t, noise):
        """ 拡散過程の一般項 """
        x_t = (
            self.sqrtab[n_t, None, None, None] * x
            + self.sqrtmab[n_t, None, None, None] * noise
        )
        return x_t

    def reverse_process(self, x_t, n_t, eps, noise):
        """ 逆拡散過程 """
        x_t = (
            self.oneover_sqrta[n_t] * (x_t - eps * self.mab_over_sqrtmab[n_t])
            + self.sqrt_beta_t[n_t] * noise
        )
        return x_t

    def forward(self, x, c):
        """ 損失関数の値を得る """
        _ts = torch.randint(1, self.n_T+1, (x.shape[0],)).to(self.device)  # 一様分布Uniform(0, n_T)に従う
        noise = torch.randn_like(x)  # 標準正規分布N(0, 1)に従う

        x_t = self.forward_process(x, _ts, noise)  

        # 条件付けの情報をマスクするかどうかはベルヌーイ分布に従う
        context_mask = torch.bernoulli(torch.zeros_like(c)+self.drop_prob).to(self.device)
        return self.loss_mse(noise, self.nn_model(x_t, c, _ts / self.n_T, context_mask))

    def classifier_free_guidance(self, eps_condition, eps_noncondition, guide_w):
        """ Classifier Free Guidance (CFG) で制御された予測値を得る """
        eps = (1+guide_w) * eps_condition - guide_w * eps_noncondition
        return eps

    def sample(self, n_sample, size, device, guide_w = 0.0):
        """ 全てのクラスで画像を生成し, その途中経過も得る 
        注意として生成個数n_sampleはクラスラベルの個数で割り切れる値になるようにする
        """
        x_i = torch.randn(n_sample, *size).to(device)  # ノイズ付き画像の初期値は標準正規分布N(0, 1)からサンプリングする
        c_i = torch.arange(0,10).to(device) # クラスラベルを一通り用意する
        c_i = c_i.repeat(int(n_sample/c_i.shape[0])) # クラスラベルを複製する. 

        # 生成個数分のコンテキストマスクの作成
        context_mask = torch.zeros_like(c_i).to(device)

        # CFGのために個数を倍増させる
        c_i = c_i.repeat(2)
        context_mask = context_mask.repeat(2)
        context_mask[n_sample:] = 1. # 後ろ半分は条件無し生成に指定する

        x_i_store = [] # 生成途中のノイズ付き画像を保存する
        for i in range(self.n_T, 0, -1):
            print(f'sampling timestep {i}',end='\r')
            # 時間ステップ数をmin-max正規化する
            t_is = torch.tensor([i / self.n_T]).to(device)
            t_is = t_is.repeat(n_sample,1,1,1) # 複製する

            # CFGのために個数を倍増させる
            x_i = x_i.repeat(2,1,1,1)
            t_is = t_is.repeat(2,1,1,1)
            
            # 逆拡散過程に用いるノイズ. 最終ステップではノイズ無し
            z = torch.randn(n_sample, *size).to(device) if i > 1 else 0

            # split predictions and compute weighting
            eps = self.nn_model(x_i, c_i, t_is, context_mask)
            eps1 = eps[:n_sample] # 条件付き生成タスクでの予測値
            eps2 = eps[n_sample:] # 条件無し生成タスクでの予測値
 
            eps = self.classifier_free_guidance(eps1, eps2, guide_w)
            x_i = x_i[:n_sample] # 個数を倍増させた分元に戻す
            x_i = self.reverse_process(x_i, i, eps, z)
            if i%20==0 or i==self.n_T or i<8:
                x_i_store.append(x_i.detach().cpu().numpy()) # 生成途中の画像を保存
        
        x_i_store = np.array(x_i_store) 
        return x_i, x_i_store

def train_mnist():
    """ メイン関数 """
    # 設定
    n_epoch = 20 # エポック数, 20エポック回すとある程度学習する
    batch_size = 256 # バッチサイズ
    n_T = 400 # 総時間ステップ数
    device = "cuda:0" # Pytorchのデバイス設定
    n_classes = 10 # クラスラベル数
    n_feat = 128 # ネットワークの特徴量の次元
    lrate = 1e-4 # 学習率の初期値
    save_model = False # 最後にモデルを保存するかどうか
    save_dir = './data/diffusion_outputs_test/' # 保存先のフォルダ
    ws_test = [0.0, 0.5, 2.0] # 条件付きへの忠実度の制御する変数のリスト
    
    # 拡散モデル
    ddpm = DDPM(nn_model=ContextUnet(in_channels=1, n_feat=n_feat, n_classes=n_classes), betas=(1e-4, 0.02), n_T=n_T, device=device, drop_prob=0.1)
    ddpm.to(device)

    # 学習済みモデルを読み込む場合
    # ddpm.load_state_dict(torch.load("PATH"))

    tf = transforms.Compose([transforms.ToTensor()]) # mnist is already normalised 0 to 1

    dataset = MNIST("./data", train=True, download=True, transform=tf) # データセット
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=5) # データローダー
    optim = torch.optim.Adam(ddpm.parameters(), lr=lrate) # 最適化アルゴリズム Adam

    for ep in range(n_epoch):
        print(f'epoch {ep}')
        ddpm.train() # 学習モード

        # 線形的に学習率を下げる
        optim.param_groups[0]['lr'] = lrate*(1-ep/n_epoch) 

        pbar = tqdm(dataloader)
        loss_ema = None # 損失関数の移動平均
        for x, c in pbar:
            optim.zero_grad()
            x = x.to(device) # 画像
            c = c.to(device) # クラスラベル
            loss = ddpm(x, c) # DDPM.forward メソッド
            loss.backward() # 勾配の計算
            if loss_ema is None:
                loss_ema = loss.item()
            else:
                loss_ema = 0.95 * loss_ema + 0.05 * loss.item()
            pbar.set_description(f"loss: {loss_ema:.4f}")
            optim.step() # パラメータの更新
        
        # 評価のため図を描画する
        ddpm.eval() # 評価モード
        with torch.no_grad():
            # 生成画像（図上段）と実画像（図下段）を比較する図を描画する
            n_sample = 4*n_classes # 生成個数
            for w_i, w in enumerate(ws_test):
                x_gen, x_gen_store = ddpm.sample(n_sample, (1, 28, 28), device, guide_w=w)

                # 学習データのバッチから指定したクラスの実画像をあるだけ集める
                x_real = torch.Tensor(x_gen.shape).to(device)
                for k in range(n_classes):
                    for j in range(int(n_sample/n_classes)):
                        try: 
                            idx = torch.squeeze((c == k).nonzero())[j]
                        except: # 指定したクラスの実画像が集めきれなかった場合
                            idx = 0
                        x_real[k+(j*n_classes)] = x[idx]

                x_all = torch.cat([x_gen, x_real]) # 生成画像と実画像の結合
                grid = make_grid(x_all*-1 + 1, nrow=10) # 画像をグリッド状に並べる
                save_image(grid, save_dir + f"image_ep{ep}_w{w}.png") # 保存
                print('saved image at ' + save_dir + f"image_ep{ep}_w{w}.png")

                if ep%5==0 or ep == int(n_epoch-1):
                    # 画像の生成の経過をgifアニメで描く
                    fig, axs = plt.subplots(nrows=int(n_sample/n_classes), ncols=n_classes,sharex=True,sharey=True,figsize=(8,3))
                    def animate_diff(i, x_gen_store):
                        """ アニメの１コマ を描く """
                        print(f'gif animating frame {i} of {x_gen_store.shape[0]}', end='\r')
                        plots = []
                        for row in range(int(n_sample/n_classes)):
                            for col in range(n_classes):
                                axs[row, col].clear()
                                axs[row, col].set_xticks([])
                                axs[row, col].set_yticks([])
                                # plots.append(axs[row, col].imshow(x_gen_store[i,(row*n_classes)+col,0],cmap='gray'))
                                plots.append(axs[row, col].imshow(-x_gen_store[i,(row*n_classes)+col,0],cmap='gray',vmin=(-x_gen_store[i]).min(), vmax=(-x_gen_store[i]).max()))
                        return plots
                    ani = FuncAnimation(fig, animate_diff, fargs=[x_gen_store],  interval=200, blit=False, repeat=True, frames=x_gen_store.shape[0])    
                    ani.save(save_dir + f"gif_ep{ep}_w{w}.gif", dpi=100, writer=PillowWriter(fps=5)) # 保存
                    print('saved image at ' + save_dir + f"gif_ep{ep}_w{w}.gif")
        # モデルを保存する
        if save_model and ep == int(n_epoch-1):
            torch.save(ddpm.state_dict(), save_dir + f"model_{ep}.pth")
            print('saved model at ' + save_dir + f"model_{ep}.pth")

if __name__ == "__main__":
    train_mnist()

