import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, random_split, SubsetRandomSampler
from tqdm import tqdm

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

transforms_train_modern = transforms.Compose([
    transforms.Pad(2),
    transforms.RandomRotation(10),
    transforms.RandomAffine(0, translate=(0.1, 0.1)),
    transforms.ToTensor(),
    transforms.Normalize(0.5, 0.5)
])

transforms_test_modern = transforms.Compose([
    transforms.Pad(2),
    transforms.ToTensor(),
    transforms.Normalize(0.5, 0.5)
])

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

train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=4, pin_memory=True, persistent_workers=True)
val_loader = DataLoader(test_dataset, batch_size=128, shuffle=False, num_workers=4, pin_memory=True, persistent_workers=True)

class LeNetModern(nn.Module):
  def __init__(self):
    super(LeNetModern, self).__init__()
    self.features = nn.Sequential(
        nn.Conv2d(1, 6, kernel_size=5, stride=1),
        nn.BatchNorm2d(6),
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=2, stride=2),

        nn.Conv2d(6, 16, kernel_size=5, stride=1),
        nn.BatchNorm2d(16),
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=2, stride=2)
    )

    self.classifier = nn.Sequential(
        nn.Linear(16 * 5 * 5, 120),
        nn.BatchNorm1d(120),
        nn.ReLU(),

        nn.Linear(120, 84),
        nn.BatchNorm1d(84),
        nn.ReLU(),

        nn.Linear(84, 10)
    )

    self._initialize_weights()

  def forward(self, x):
    x = self.features(x)
    x = torch.flatten(x, 1)
    x = self.classifier(x)
    return x

  def _initialize_weights(self):
    for m in self.modules():
      if isinstance(m, nn.Conv2d):
        nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
        if m.bias is not None:
          nn.init.constant_(m.bias, 0)
      if isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.BatchNorm1d):
        nn.init.constant_(m.weight, 1)
        nn.init.constant_(m.bias, 0)
      if isinstance(m, nn.Linear):
        nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity='relu')
        if m.bias is not None:
          nn.init.constant_(m.bias, 0)

model = LeNetModern().to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)

best_acc = 0
epochs = 50

for epoch in range(epochs):
  model.train()
  running_loss = 0.0
  correct = 0
  total = 0

  train_pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs} [Train]")
  for inputs, targets in train_pbar:
    inputs, targets = inputs.to(device), targets.to(device)
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()

    running_loss += loss.item()
    _, predicted = outputs.max(1)
    total += targets.size(0)
    correct += predicted.eq(targets).sum().item()

    train_pbar.set_postfix(loss=running_loss/(total/targets.size(0)), acc=100.*correct/total)

  train_loss = running_loss / len(train_loader)
  train_acc = 100. * correct / total

  model.eval()
  val_loss = 0.0
  val_correct = 0
  val_total = 0

  with torch.no_grad():
    val_pbar = tqdm(val_loader, desc=f"Epoch {epoch+1}/{epochs} [Val]")
    for inputs, targets in val_pbar:
      inputs, targets = inputs.to(device), targets.to(device)
      outputs = model(inputs)
      loss = criterion(outputs, targets)

      val_loss += loss.item()
      _, predicted = outputs.max(1)
      val_total += targets.size(0)
      val_correct += predicted.eq(targets).sum().item()
      val_pbar.set_postfix(loss=val_loss/(val_total/targets.size(0)), acc=100.*val_correct/val_total)

  validation_loss = val_loss / len(val_loader)
  validation_acc = 100. * val_correct / val_total

  print(f"Epoch {epoch+1}/{epochs} | Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}% | Val Loss: {validation_loss:.4f}, Val Acc: {validation_acc:.2f}%")

  scheduler.step()

  if validation_acc > best_acc:
    best_acc = validation_acc
    torch.save(model.state_dict(), 'best_lenet_classic_model.pth')
    print(f"New best model found! Accuracy: {best_acc:.2f}%. Saving model...")

print("\nTraining finished.")
print(f"Best validation accuracy: {best_acc:.2f}%")

Epoch 1/50 [Train]: 100%|██████████| 469/469 [00:07<00:00, 60.12it/s, acc=71.5, loss=0.615]
Epoch 1/50 [Val]: 100%|██████████| 79/79 [00:00<00:00, 89.18it/s, acc=80, loss=0.0665]

Epoch 1/50 | Train Loss: 0.8200, Train Acc: 71.52% | Val Loss: 0.5262, Val Acc: 80.01%
New best model found! Accuracy: 80.01%. Saving model...

Epoch 2/50 [Train]: 100%|██████████| 469/469 [00:06<00:00, 70.48it/s, acc=79.8, loss=0.406]
Epoch 2/50 [Val]: 100%|██████████| 79/79 [00:00<00:00, 118.25it/s, acc=84.1, loss=0.0559]

Epoch 2/50 | Train Loss: 0.5413, Train Acc: 79.84% | Val Loss: 0.4419, Val Acc: 84.10%
New best model found! Accuracy: 84.10%. Saving model...

Epoch 3/50 [Train]: 100%|██████████| 469/469 [00:06<00:00, 71.22it/s, acc=82.2, loss=0.358]
Epoch 3/50 [Val]: 100%|██████████| 79/79 [00:00<00:00, 117.15it/s, acc=86, loss=0.049]

Epoch 3/50 | Train Loss: 0.4775, Train Acc: 82.21% | Val Loss: 0.3877, Val Acc: 86.02%
New best model found! Accuracy: 86.02%. Saving model...

Epoch 4/50 [Train]: 100%|██████████| 469/469 [00:06<00:00, 70.07it/s, acc=83.6, loss=0.331]
Epoch 4/50 [Val]: 100%|██████████| 79/79 [00:00<00:00, 114.84it/s, acc=86.3, loss=0.0465]

Epoch 4/50 | Train Loss: 0.4416, Train Acc: 83.58% | Val Loss: 0.3682, Val Acc: 86.34%
New best model found! Accuracy: 86.34%. Saving model...

Epoch 5/50 [Train]: 100%|██████████| 469/469 [00:06<00:00, 71.37it/s, acc=84.8, loss=0.307]
Epoch 5/50 [Val]: 100%|██████████| 79/79 [00:00<00:00, 107.97it/s, acc=86.3, loss=0.0463]

Epoch 5/50 | Train Loss: 0.4085, Train Acc: 84.82% | Val Loss: 0.3661, Val Acc: 86.29%