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, SubsetRandomSampler
from tqdm import tqdm

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

CIFAR10_MEAN = [0.4914, 0.4822, 0.4465]
CIFAR10_STD = [0.2470, 0.2435, 0.2616]

train_transform = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(CIFAR10_MEAN, CIFAR10_STD),
])

test_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(CIFAR10_MEAN, CIFAR10_STD)
])

train_dataset = datasets.CIFAR10(
    root='./data',
    train=True,
    transform=train_transform,
    download=True
)

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

train_loader = DataLoader(
    dataset=train_dataset,
    batch_size=128,
    shuffle=True,
    num_workers=4,
    pin_memory=True,
    persistent_workers=True
)

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

class NiN(nn.Module):
  def __init__(self, num_classes=10):
    super(NiN, self).__init__()
    self.features = nn.Sequential(
        # MLE Block 1
        nn.Conv2d(3, 192, kernel_size=5, stride=1, padding=2),
        nn.BatchNorm2d(192),
        nn.ReLU(),
        nn.Conv2d(192, 160, kernel_size=1, stride=1, padding=0),
        nn.BatchNorm2d(160),
        nn.ReLU(),
        nn.Conv2d(160, 96, kernel_size=1, stride=1, padding=0),
        nn.BatchNorm2d(96),
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
        nn.Dropout(0.2),

        nn.Conv2d(96, 192, kernel_size=5, stride=1, padding=2),
        nn.BatchNorm2d(192),
        nn.ReLU(),
        nn.Conv2d(192, 192, kernel_size=1, stride=1, padding=0),
        nn.BatchNorm2d(192),
        nn.ReLU(),
        nn.Conv2d(192, 192, kernel_size=1, stride=1, padding=0),
        nn.BatchNorm2d(192),
        nn.ReLU(),
        nn.AvgPool2d(kernel_size=3, stride=2, padding=1),
        nn.Dropout(0.2),

        nn.Conv2d(192, 192, kernel_size=3, stride=1, padding=1),
        nn.BatchNorm2d(192),
        nn.ReLU(),
        nn.Conv2d(192, 192, kernel_size=1, stride=1, padding=0),
        nn.BatchNorm2d(192),
        nn.ReLU(),
        nn.Conv2d(192, num_classes, kernel_size=1, stride=1, padding=0),

        nn.AdaptiveAvgPool2d((1, 1))
    )

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

model = NiN().to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[30, 60, 80], gamma=0.2)

best_acc = 0
epochs = 100

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/100 [Train]: 100%|██████████| 391/391 [00:13<00:00, 28.27it/s, acc=47.6, loss=0.892]
Epoch 1/100 [Val]: 100%|██████████| 79/79 [00:01<00:00, 73.69it/s, acc=49.4, loss=0.192]

Epoch 1/100 | Train Loss: 1.4260, Train Acc: 47.60% | Val Loss: 1.5218, Val Acc: 49.36%
New best model found! Accuracy: 49.36%. Saving model...

Epoch 2/100 [Train]: 100%|██████████| 391/391 [00:13<00:00, 28.29it/s, acc=62.8, loss=0.651]
Epoch 2/100 [Val]: 100%|██████████| 79/79 [00:01<00:00, 78.96it/s, acc=63.8, loss=0.131]

Epoch 2/100 | Train Loss: 1.0411, Train Acc: 62.81% | Val Loss: 1.0336, Val Acc: 63.78%
New best model found! Accuracy: 63.78%. Saving model...

Epoch 3/100 [Train]: 100%|██████████| 391/391 [00:13<00:00, 28.31it/s, acc=69.9, loss=0.539]
Epoch 3/100 [Val]: 100%|██████████| 79/79 [00:00<00:00, 79.21it/s, acc=60.7, loss=0.161]

Epoch 3/100 | Train Loss: 0.8617, Train Acc: 69.90% | Val Loss: 1.2754, Val Acc: 60.69%

Epoch 4/100 [Train]: 100%|██████████| 391/391 [00:13<00:00, 28.45it/s, acc=73.1, loss=0.483]
Epoch 4/100 [Val]: 100%|██████████| 79/79 [00:00<00:00, 79.70it/s, acc=68.4, loss=0.12]

Epoch 4/100 | Train Loss: 0.7718, Train Acc: 73.09% | Val Loss: 0.9493, Val Acc: 68.37%
New best model found! Accuracy: 68.37%. Saving model...

Epoch 5/100 [Train]: 100%|██████████| 391/391 [00:13<00:00, 28.54it/s, acc=75.7, loss=0.442]
Epoch 5/100 [Val]: 100%|██████████| 79/79 [00:00<00:00, 79.67it/s, acc=73.2, loss=0.0999]

Epoch 5/100 | Train Loss: 0.7073, Train Acc: 75.65% | Val Loss: 0.7903, Val Acc: 73.24%
New best model found! Accuracy: 73.24%. Saving model...