###########################################################################################
# Copyright (c) Hayeon Lee, Eunyoung Hyung [GitHub MetaD2A], 2021
# Rapid Neural Architecture Search by Learning to Generate Graphs from Datasets, ICLR 2021
###########################################################################################
import argparse

def str2bool(v):
  return v.lower() in ['t', 'true', True]


def get_parser():
  parser = argparse.ArgumentParser()
  # general settings
  parser.add_argument('--seed', type=int, default=333)
  parser.add_argument('--gpu', type=str, default='0', help='set visible gpus')
  parser.add_argument('--model_name', type=str, default=None, choices=['generator', 'predictor', 'train_arch'])
  parser.add_argument('--save-path', type=str, default='results', help='the path of save directory')
  parser.add_argument('--data-path', type=str, default='data', help='the path of save directory')
  parser.add_argument('--save-epoch', type=int, default=20, help='how many epochs to wait each time to save model states') 
  parser.add_argument('--max-epoch', type=int, default=400, help='number of epochs to train')
  parser.add_argument('--batch_size', type=int, default=32, help='batch size for generator')
  parser.add_argument('--graph-data-name', default='ofa_mbv3', help='graph dataset name')
  parser.add_argument('--nvt', type=int, default=27, help='number of different node types, 21 for ofa_mbv3 without in/out node')
  # set encoder
  parser.add_argument('--num-sample', type=int, default=20, help='the number of images as input for set encoder')
  # graph encoder
  parser.add_argument('--hs', type=int, default=56, help='hidden size of GRUs')
  parser.add_argument('--nz', type=int, default=56, help='the number of dimensions of latent vectors z')
  # test
  parser.add_argument('--test', action='store_true', default=False, help='turn on test mode')
  parser.add_argument('--load-epoch', type=int, default=20, help='checkpoint epoch loaded for meta-test') 
  parser.add_argument('--data-name', type=str, default=None, help='meta-test dataset name') 
  parser.add_argument('--num-class', type=int, default=None, help='the number of class of dataset') 
  parser.add_argument('--num-gen-arch', type=int, default=200, help='the number of candidate architectures generated by the generator') 
  parser.add_argument('--train-arch', type=str2bool, default=True, help='whether to train the searched architecture') 
  # database
  parser.add_argument('--index', type=int, default=None, help='the process number when creating DB')
  parser.add_argument('--imgnet', type=str, default=None, help='The path of imagenet')
  parser.add_argument('--collect', action='store_true', default=False, help='whether to train the searched architecture')

  args = parser.parse_args()
  
  return args