import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# デバイス設定
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# --- データの準備 ---
transform = transforms.Compose([
    transforms.ToTensor(),
])

train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=256, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=256, shuffle=False)

# --- オートエンコーダの定義 ---
class AutoEncoder(nn.Module):
    def __init__(self, input_dim=784, latent_dim=2):
        super().__init__()
        # エンコーダ
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 64),
            nn.ReLU(),
            nn.Linear(64, latent_dim),
        )
        # デコーダ
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 256),
            nn.ReLU(),
            nn.Linear(256, input_dim),
            nn.Sigmoid(),
        )

    def forward(self, x):
        z = self.encoder(x)
        x_hat = self.decoder(z)
        return x_hat, z

# --- 学習 ---
model = AutoEncoder(input_dim=784, latent_dim=2).to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.MSELoss()

n_epochs = 20
train_losses = []

for epoch in range(n_epochs):
    model.train()
    epoch_loss = 0
    for batch_x, _ in train_loader:
        batch_x = batch_x.view(-1, 784).to(device)
        x_hat, z = model(batch_x)
        loss = criterion(x_hat, batch_x)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()

    avg_loss = epoch_loss / len(train_loader)
    train_losses.append(avg_loss)
    if (epoch + 1) % 5 == 0:
        print(f"Epoch [{epoch+1}/{n_epochs}], Loss: {avg_loss:.6f}")

# --- 可視化 ---
# 学習曲線
fig, ax = plt.subplots(figsize=(8, 4))
ax.plot(train_losses, 'b-')
ax.set_xlabel('Epoch')
ax.set_ylabel('MSE Loss')
ax.set_title('Training Loss')
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

# 再構成結果
model.eval()
with torch.no_grad():
    test_x, test_y = next(iter(test_loader))
    test_x_flat = test_x.view(-1, 784).to(device)
    recon, latent = model(test_x_flat)

fig, axes = plt.subplots(2, 10, figsize=(15, 3))
for i in range(10):
    axes[0, i].imshow(test_x[i].squeeze(), cmap='gray')
    axes[0, i].axis('off')
    axes[1, i].imshow(recon[i].cpu().view(28, 28), cmap='gray')
    axes[1, i].axis('off')
axes[0, 0].set_ylabel('Original')
axes[1, 0].set_ylabel('Reconstructed')
plt.suptitle('AutoEncoder Reconstruction', fontsize=14)
plt.tight_layout()
plt.show()

# 潜在空間の可視化（2次元）
model.eval()
all_z = []
all_labels = []
with torch.no_grad():
    for batch_x, batch_y in test_loader:
        batch_x = batch_x.view(-1, 784).to(device)
        _, z = model(batch_x)
        all_z.append(z.cpu().numpy())
        all_labels.append(batch_y.numpy())

all_z = np.concatenate(all_z)
all_labels = np.concatenate(all_labels)

fig, ax = plt.subplots(figsize=(8, 8))
scatter = ax.scatter(all_z[:, 0], all_z[:, 1], c=all_labels,
                     cmap='tab10', s=5, alpha=0.5)
plt.colorbar(scatter, ax=ax, label='Digit')
ax.set_xlabel('Latent dim 1')
ax.set_ylabel('Latent dim 2')
ax.set_title('Latent Space Visualization (2D)')
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()