import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

# データを取得(引数の詳細：https://pytorch.org/vision/main/generated/torchvision.datasets.MNIST.html)
data = MNIST("./data",
             train=True,
             download=True,
             transform=transforms.ToTensor())

# 訓練データとテストデータの割合
train_size = int(len(data) * 0.8)
test_size = int(len(data) * 0.2)
train_data, test_data = torch.utils.data.random_split(data, [train_size, test_size])

# バッチごとにデータを読み込む(バッチサイズ100)
train_loader = DataLoader(dataset=train_data,
                          batch_size=100,
                          shuffle=True,
                          num_workers=0)

test_loader = DataLoader(dataset=test_data,
                        batch_size=100,
                        shuffle=True,
                        num_workers=0)

# 最初のバッチのデータを取り出す
images, labels = next(iter(train_loader))
print("images_size:",images.size())   #images_size: torch.Size([100, 1, 28, 28])

z_dim = 2
num_epochs = 20

# おまじない(詳細：https://atmarkit.itmedia.co.jp/ait/articles/2008/28/news030.html)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# インスタンス化
model = VAE(z_dim).to(device)
# 最適化アルゴリズム（今回はAdam使用）
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[15], gamma=0.1)

# loss、分散、平均など記録用
history = {"train_loss": [], "test_loss": [], "ave": [], "log_dev": [], "z": [], "labels":[]}

for epoch in range(num_epochs):
    model.train() # 学習モーード☆に切り替え

    # 訓練データを取り出し学習
    for i, (x, labels) in enumerate(train_loader):
        # 入力xからxを再構成
        input = x.to(device).view(-1, 28*28).to(torch.float32)
        output, z, ave, log_dev = model(input)

        # 諸々記録
        history["ave"].append(ave)
        history["log_dev"].append(log_dev)
        history["z"].append(z)
        history["labels"].append(labels)
        loss = criterion(output, input, ave, log_dev)

        # 最適化、誤差逆伝播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (i+1) % 50 == 0:
            print(f'Epoch: {epoch+1}, loss: {loss: 0.4f}')
        history["train_loss"].append(loss)

    model.eval() # 推論ターイム☆☆に切り替え
    with torch.no_grad():
        # 訓練データを取り出し学習
        for i, (x, labels) in enumerate(test_loader):
            # 入力xからxを再構成
            input = x.to(device).view(-1, 28*28).to(torch.float32)
            output, z, ave, log_dev = model(input)

            loss = criterion(output, input, ave, log_dev)
            # lossが下がっているか(うまく学習できているか)の確認のため記録
            history["test_loss"].append(loss)

        print(f'Epoch: {epoch+1}, test_loss: {loss: 0.4f}')

    scheduler.step()

# numpyに変換してからplot
train_loss_tensor = torch.stack(history["test_loss"]) # ここをtrain_lossにすることで訓練データのlossも見れる
train_loss_np = train_loss_tensor.detach().numpy().copy()
plt.plot(train_loss_np)

# zと正解ラベルをnumpyに変換
z_tensor = torch.stack(history["z"])
z_np = z_tensor.detach().numpy().copy()
labels_tensor = torch.stack(history["labels"])
labels_np = labels_tensor.detach().numpy().copy()

cmap_keyword = "tab10"
cmap = plt.get_cmap(cmap_keyword)

# バッチが9550以降のzを可視化
batch_num = 9550
plt.figure(figsize=[10,10])
for label in range(10):
    x = z_np[batch_num:,:,0][labels_np[batch_num:,:] == label]
    y = z_np[batch_num:,:,1][labels_np[batch_num:,:] == label]
    plt.scatter(x, y, color=cmap(label/9), label=label, s=15)
    plt.annotate(label, xy=(np.mean(x),np.mean(y)),size=20,color="black")
plt.legend(loc="upper left")

f, ax = plt.subplots(10, 10, figsize=(8, 8))

for i, x in enumerate(np.linspace(-1.5, 1.5, 10)):
    for j, y in enumerate(np.linspace(-2, 2, 10)):
        # エンコーダにzを入れた出力をplot
        z = torch.tensor([x, y], dtype = torch.float32)
        output = model.decoder(z)
        np_output = output.detach().numpy().copy()
        np_image = np.reshape(np_output, (28, 28))
        ax[9-j][i].imshow(np_image, cmap='gray')