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')
print(f"Device: {device}")

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

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

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

# --- CNNモデルの定義 ---
class CNN(nn.Module):
    def __init__(self):
        super().__init__()
        # 畳み込み層
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)   # 28x28 -> 28x28
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)  # 14x14 -> 14x14
        self.pool = nn.MaxPool2d(2, 2)                             # サイズを半分に
        self.dropout1 = nn.Dropout2d(0.25)
        self.dropout2 = nn.Dropout(0.5)
        # 全結合層
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        # Conv1 -> ReLU -> Pool: (1, 28, 28) -> (32, 14, 14)
        x = self.pool(torch.relu(self.conv1(x)))
        # Conv2 -> ReLU -> Pool: (32, 14, 14) -> (64, 7, 7)
        x = self.pool(torch.relu(self.conv2(x)))
        x = self.dropout1(x)
        # Flatten
        x = x.view(-1, 64 * 7 * 7)
        # FC1 -> ReLU -> Dropout
        x = torch.relu(self.fc1(x))
        x = self.dropout2(x)
        # FC2 (出力)
        x = self.fc2(x)
        return x

model = CNN().to(device)

# パラメータ数の確認
total_params = sum(p.numel() for p in model.parameters())
print(f"総パラメータ数: {total_params:,}")

# --- 学習 ---
optimizer = optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

n_epochs = 10
train_losses = []
test_accuracies = []

for epoch in range(n_epochs):
    model.train()
    epoch_loss = 0
    for batch_x, batch_y in train_loader:
        batch_x, batch_y = batch_x.to(device), batch_y.to(device)
        optimizer.zero_grad()
        output = model(batch_x)
        loss = criterion(output, batch_y)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()

    avg_loss = epoch_loss / len(train_loader)
    train_losses.append(avg_loss)

    # テスト精度
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for batch_x, batch_y in test_loader:
            batch_x, batch_y = batch_x.to(device), batch_y.to(device)
            output = model(batch_x)
            pred = output.argmax(dim=1)
            correct += (pred == batch_y).sum().item()
            total += batch_y.size(0)
    acc = correct / total
    test_accuracies.append(acc)
    print(f"Epoch [{epoch+1}/{n_epochs}] Loss: {avg_loss:.4f}, Test Acc: {acc:.4f}")

# --- 可視化 ---
fig, axes = plt.subplots(1, 2, figsize=(12, 4))

ax1 = axes[0]
ax1.plot(train_losses, 'b-o', markersize=4)
ax1.set_xlabel('Epoch')
ax1.set_ylabel('Loss')
ax1.set_title('Training Loss')
ax1.grid(True, alpha=0.3)

ax2 = axes[1]
ax2.plot(test_accuracies, 'r-o', markersize=4)
ax2.set_xlabel('Epoch')
ax2.set_ylabel('Accuracy')
ax2.set_title('Test Accuracy')
ax2.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

# --- 予測結果の表示 ---
model.eval()
test_x, test_y = next(iter(test_loader))
with torch.no_grad():
    output = model(test_x.to(device))
    preds = output.argmax(dim=1).cpu()

fig, axes = plt.subplots(2, 5, figsize=(12, 5))
for i in range(10):
    ax = axes[i // 5, i % 5]
    ax.imshow(test_x[i].squeeze(), cmap='gray')
    color = 'green' if preds[i] == test_y[i] else 'red'
    ax.set_title(f'Pred: {preds[i].item()} (True: {test_y[i].item()})', color=color)
    ax.axis('off')
plt.tight_layout()
plt.show()