# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the \"License\");
# you may not use this file except in compliance with the License.\n",
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an \"AS IS\" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Train resnet on single GPU."""

import argparse
import time

from PIL import Image
import torch
from torch import nn
import torchmetrics
import torchvision
from torchvision.models import resnet50


class ImageFolder(torchvision.datasets.ImageFolder):
  """Class for loading imagenet."""

  def __init__(self, image_list_file, transform=None, target_transform=None):
    self.samples = self._make_dataset(image_list_file)
    self.loader = self._loader

    self.imgs = self.samples
    self.targets = [s[1] for s in self.samples]

    self.transform = transform
    self.target_transform = target_transform

  def _make_dataset(self, image_list_file):
    items = []
    with open(image_list_file, 'r') as f:
      for line in f:
        item = line.strip().split(' ')
        items.append((item[0], int(item[1])))
    return items

  def _loader(self, image_path):
    with open(image_path, 'rb') as f:
      img = Image.open(f)
      img = img.convert('RGB')
      return img


def train(model, device, dataloader, optimizer):
  model.train()
  for image, target in dataloader:
    image, target = image.to(device), target.to(device)
    pred = model(image)
    # pred.shape (N, C), target.shape (N)
    loss = nn.functional.cross_entropy(pred, target)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
  return loss


def evaluate(model, device, dataloader, metric):
  model.eval()
  with torch.no_grad():
    for image, target in dataloader:
      image, target = image.to(device), target.to(device)
      pred = model(image)
      metric.update(pred, target)
  accuracy = metric.compute()
  metric.reset()
  return accuracy


def run_training(args):
  """Run training and evaluation."""
  # Create model.
  model = resnet50(weights=None)
  model = model.to(args.device)

  # Create train dataloader.
  train_dataset = ImageFolder(
      image_list_file=args.train_data_path,
      transform=torchvision.transforms.Compose([
          torchvision.transforms.RandomResizedCrop(224),
          torchvision.transforms.RandomHorizontalFlip(),
          torchvision.transforms.ToTensor(),
          torchvision.transforms.Normalize(
              mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
      ]))
  train_dataloader = torch.utils.data.DataLoader(
      dataset=train_dataset,
      batch_size=args.train_batch_size,
      shuffle=True,
      num_workers=args.dataloader_num_workers,
      pin_memory=True)
  print(f'Train dataloader | samples: {len(train_dataloader.dataset)}, '
        f'num workers: {train_dataloader.num_workers}, '
        f'batch size: {args.train_batch_size}, '
        f'batches/epoch: {len(train_dataloader)}')

  # Create eval dataloader.
  eval_dataset = ImageFolder(
      image_list_file=args.eval_data_path,
      transform=torchvision.transforms.Compose([
          torchvision.transforms.Resize(256),
          torchvision.transforms.CenterCrop(224),
          torchvision.transforms.ToTensor(),
          torchvision.transforms.Normalize(
              mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
      ]))
  eval_dataloader = torch.utils.data.DataLoader(
      dataset=eval_dataset,
      batch_size=args.eval_batch_size,
      shuffle=False,
      num_workers=args.dataloader_num_workers,
      pin_memory=True,
      drop_last=True)
  print(f'Eval dataloader | samples: {len(eval_dataloader.dataset)}, '
        f'num workers: {eval_dataloader.num_workers}, '
        f'batch size: {args.eval_batch_size}, '
        f'batches/epoch: {len(eval_dataloader)}')

  # Optimizer.
  optimizer = torch.optim.SGD(model.parameters(), 0.1)

  # Main loop.
  metric = torchmetrics.classification.Accuracy(top_k=1).to(args.device)
  for epoch in range(1, args.epochs + 1):
    print(f'Running epoch {epoch}')

    start = time.time()
    train(model, args.device, train_dataloader, optimizer)
    end = time.time()
    print(f'Training finished in {(end - start):>0.3f} seconds')

    start = time.time()
    evaluate(model, args.device, eval_dataloader, metric)
    end = time.time()
    print(f'Evaluation finished in {(end - start):>0.3f} seconds')
  print('Done')


def create_args():
  """Create main args."""
  parser = argparse.ArgumentParser(
      formatter_class=argparse.ArgumentDefaultsHelpFormatter)
  parser.add_argument(
      '--epochs',
      default=1,
      type=int,
      help='number of total epochs to run')
  parser.add_argument(
      '--dataloader_num_workers',
      default=2,
      type=int,
      help='number of workders for dataloader')
  parser.add_argument(
      '--train_data_path',
      default='',
      type=str,
      help='path to training data')
  parser.add_argument(
      '--train_batch_size',
      default=32,
      type=int,
      help='batch size for training')
  parser.add_argument(
      '--eval_data_path',
      default='',
      type=str,
      help='path to evaluation data')
  parser.add_argument(
      '--eval_batch_size',
      default=32,
      type=int,
      help='batch size for evaluation')
  args = parser.parse_args()
  return args


def main():
  args = create_args()
  args.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

  print('Launch job on 1 GPU')
  run_training(args)


if __name__ == '__main__':
  main()