import os
import argparse
import torch
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms

from model import Net
from azureml.core import Run

run = Run.get_context()

if __name__ == "__main__":
  parser = argparse.ArgumentParser()
  parser.add_argument(
      '--data_path',
      type=str,
      help='Path to the training data'
  )
  parser.add_argument(
      '--learning_rate',
      type=float,
      default=0.001,
      help='Learning rate for SGD'
  )
  parser.add_argument(
      '--momentum',
      type=float,
      default=0.9,
      help='Momentum for SGD'
  )

  args = parser.parse_args()

  print("===== DATA =====")
  print("DATA PATH: " + args.data_path)
  print("LIST FILES IN DATA PATH...")
  print(os.listdir(args.data_path))
  print("================")

  # prepare DataLoader for CIFAR10 data
  transform = transforms.Compose([
      transforms.ToTensor(),
      transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
  ])
  trainset = torchvision.datasets.CIFAR10(
      root=args.data_path,
      train=True,
      download=False,
      transform=transform,
  )
  trainloader = torch.utils.data.DataLoader(
      trainset,
      batch_size=4,
      shuffle=True,
      num_workers=2
  )

  # define convolutional network
  net = Net()

  # set up pytorch loss /  optimizer
  criterion = torch.nn.CrossEntropyLoss()
  optimizer = optim.SGD(
      net.parameters(),
      lr=args.learning_rate,
      momentum=args.momentum,
  )

  # train the network
  for epoch in range(2):

    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
      # unpack the data
      inputs, labels = data

      # zero the parameter gradients
      optimizer.zero_grad()

      # forward + backward + optimize
      outputs = net(inputs)
      loss = criterion(outputs, labels)
      loss.backward()
      optimizer.step()

      # print statistics
      running_loss += loss.item()
      if i % 2000 == 1999:
        loss = running_loss / 2000
        run.log('loss', loss)  # log loss metric to AML
        print(f'epoch={epoch + 1}, batch={i + 1:5}: loss {loss:.2f}')
        running_loss = 0.0

  print('Finished Training')
[((11, 6, 11, 23), 'azureml.core.Run.get_context', 'Run.get_context', ({}, {}), '()', False, 'from azureml.core import Run\n'), ((14, 11, 14, 36), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((46, 13, 51, 3), 'torchvision.datasets.CIFAR10', 'torchvision.datasets.CIFAR10', (), '', False, 'import torchvision\n'), ((52, 16, 57, 3), 'torch.utils.data.DataLoader', 'torch.utils.data.DataLoader', (), '', False, 'import torch\n'), ((60, 8, 60, 13), 'model.Net', 'Net', ({}, {}), '()', False, 'from model import Net\n'), ((63, 14, 63, 41), 'torch.nn.CrossEntropyLoss', 'torch.nn.CrossEntropyLoss', ({}, {}), '()', False, 'import torch\n'), ((38, 8, 38, 34), 'os.listdir', 'os.listdir', ({(38, 19, 38, 33): 'args.data_path'}, {}), '(args.data_path)', False, 'import os\n'), ((43, 6, 43, 27), 'torchvision.transforms.ToTensor', 'transforms.ToTensor', ({}, {}), '()', True, 'import torchvision.transforms as transforms\n'), ((44, 6, 44, 60), 'torchvision.transforms.Normalize', 'transforms.Normalize', ({(44, 27, 44, 42): '(0.5, 0.5, 0.5)', (44, 44, 44, 59): '(0.5, 0.5, 0.5)'}, {}), '((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))', True, 'import torchvision.transforms as transforms\n')]