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]

transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])

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 DepthwiseSeparableConv(nn.Module):
  def __init__(self, in_channels, out_channels, stride=1):
    super(DepthwiseSeparableConv, self).__init__()

    self.depthwise = nn.Conv2d(
      in_channels,
      in_channels,
      kernel_size=3,
      stride=stride,
      padding=1,
      groups=in_channels,
      bias=False
    )
    self.bn1 = nn.BatchNorm2d(in_channels)
    self.relu1 = nn.ReLU(inplace=True)

    self.pointwise = nn.Conv2d(
      in_channels,
      out_channels,
      kernel_size=1,
      stride=1,
      padding=0,
      bias=False
    )
    self.bn2 = nn.BatchNorm2d(out_channels)
    self.relu2 = nn.ReLU(inplace=True)

  def forward(self, x):
    x = self.depthwise(x)
    x = self.bn1(x)
    x = self.relu1(x)
    x = self.pointwise(x)
    x = self.bn2(x)
    x = self.relu2(x)
    return x


class MobileNetV1(nn.Module):
  def __init__(self, num_classes=10, width_multiplier=1.0):
    super(MobileNetV1, self).__init__()

    self.conv1 = nn.Sequential(
      nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1, bias=False),
      nn.BatchNorm2d(32),
      nn.ReLU(inplace=True)
    )

    self.features = nn.Sequential(
      DepthwiseSeparableConv(32, 64, stride=1),    # 32x32
      DepthwiseSeparableConv(64, 128, stride=2),   # 16x16
      DepthwiseSeparableConv(128, 128, stride=1),  # 16x16
      DepthwiseSeparableConv(128, 256, stride=2),  # 8x8
      DepthwiseSeparableConv(256, 256, stride=1),  # 8x8
      DepthwiseSeparableConv(256, 512, stride=2),  # 4x4
      DepthwiseSeparableConv(512, 512, stride=1),  # 4x4
      DepthwiseSeparableConv(512, 512, stride=1),  # 4x4
      DepthwiseSeparableConv(512, 512, stride=1),  # 4x4
    )

    self.avgpool = nn.AdaptiveAvgPool2d(1)
    self.dropout = nn.Dropout(0.2)
    self.classifier = nn.Linear(512, num_classes)

    self._initialize_weights()

  def forward(self, x):
    x = self.conv1(x)
    x = self.features(x)
    x = self.avgpool(x)
    x = torch.flatten(x, 1)
    x = self.dropout(x)
    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.zeros_(m.bias)
      elif isinstance(m, nn.BatchNorm2d):
        nn.init.ones_(m.weight)
        nn.init.zeros_(m.bias)
      elif isinstance(m, nn.Linear):
        nn.init.normal_(m.weight, 0, 0.01)
        nn.init.zeros_(m.bias)

model = MobileNetV1().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:05<00:00, 65.91it/s, acc=40, loss=1.03]
Epoch 1/100 [Val]: 100%|██████████| 79/79 [00:00<00:00, 97.53it/s, acc=55.3, loss=0.156]

Epoch 1/100 | Train Loss: 1.6462, Train Acc: 39.95% | Val Loss: 1.2343, Val Acc: 55.26%
New best model found! Accuracy: 55.26%. Saving model...

Epoch 2/100 [Train]: 100%|██████████| 391/391 [00:05<00:00, 65.88it/s, acc=60.7, loss=0.689]
Epoch 2/100 [Val]: 100%|██████████| 79/79 [00:00<00:00, 108.17it/s, acc=61.2, loss=0.143]

Epoch 2/100 | Train Loss: 1.1012, Train Acc: 60.69% | Val Loss: 1.1295, Val Acc: 61.21%
New best model found! Accuracy: 61.21%. Saving model...

Epoch 3/100 [Train]: 100%|██████████| 391/391 [00:05<00:00, 68.95it/s, acc=68.8, loss=0.552]
Epoch 3/100 [Val]: 100%|██████████| 79/79 [00:00<00:00, 111.18it/s, acc=65.2, loss=0.127]

Epoch 3/100 | Train Loss: 0.8829, Train Acc: 68.77% | Val Loss: 1.0030, Val Acc: 65.23%
New best model found! Accuracy: 65.23%. Saving model...

Epoch 4/100 [Train]: 100%|██████████| 391/391 [00:05<00:00, 66.32it/s, acc=73.8, loss=0.477]
Epoch 4/100 [Val]: 100%|██████████| 79/79 [00:00<00:00, 109.22it/s, acc=71.3, loss=0.104]

Epoch 4/100 | Train Loss: 0.7621, Train Acc: 73.80% | Val Loss: 0.8262, Val Acc: 71.27%
New best model found! Accuracy: 71.27%. Saving model...

Epoch 5/100 [Train]: 100%|██████████| 391/391 [00:05<00:00, 67.25it/s, acc=76, loss=0.432]
Epoch 5/100 [Val]: 100%|██████████| 79/79 [00:00<00:00, 101.03it/s, acc=75.3, loss=0.0912]

Epoch 5/100 | Train Loss: 0.6908, Train Acc: 76.00% | Val Loss: 0.7217, Val Acc: 75.31%
New best model found! Accuracy: 75.31%. Saving model...