import os, sys, time, random, argparse
from .share_args import add_shared_args


def obtain_search_single_args():
  parser = argparse.ArgumentParser(description='Train a classification model on typical image classification datasets.', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
  parser.add_argument('--resume'        ,   type=str,                   help='Resume path.')
  parser.add_argument('--model_config'  ,   type=str,                   help='The path to the model configuration')
  parser.add_argument('--optim_config'  ,   type=str,                   help='The path to the optimizer configuration')
  parser.add_argument('--split_path'    ,   type=str,                   help='The split file path.')
  parser.add_argument('--search_shape'  ,   type=str,                   help='The shape to be searched.')
  #parser.add_argument('--arch_para_pure',   type=int,                   help='The architecture-parameter pure or not.')
  parser.add_argument('--gumbel_tau_max',   type=float,                 help='The maximum tau for Gumbel.')
  parser.add_argument('--gumbel_tau_min',   type=float,                 help='The minimum tau for Gumbel.')
  parser.add_argument('--procedure'     ,   type=str,                   help='The procedure basic prefix.')
  parser.add_argument('--FLOP_ratio'    ,   type=float,                 help='The expected FLOP ratio.')
  parser.add_argument('--FLOP_weight'   ,   type=float,                 help='The loss weight for FLOP.')
  parser.add_argument('--FLOP_tolerant' ,   type=float,                 help='The tolerant range for FLOP.')
  add_shared_args( parser )
  # Optimization options
  parser.add_argument('--batch_size'    ,   type=int,   default=2,      help='Batch size for training.')
  args = parser.parse_args()

  if args.rand_seed is None or args.rand_seed < 0:
    args.rand_seed = random.randint(1, 100000)
  assert args.save_dir is not None, 'save-path argument can not be None'
  assert args.gumbel_tau_max is not None and args.gumbel_tau_min is not None
  assert args.FLOP_tolerant is not None and args.FLOP_tolerant > 0, 'invalid FLOP_tolerant : {:}'.format(FLOP_tolerant)
  #assert args.arch_para_pure is not None, 'arch_para_pure is not None: {:}'.format(args.arch_para_pure)
  #args.arch_para_pure = bool(args.arch_para_pure)
  return args