import argparse


def parser():

    # Argument
    parser = argparse.ArgumentParser()

    # Directory
    parser.add_argument('--data_dir', default='/data_large/readonly', type=str, help='directory to CIFAR-10 and ImageNet data.')
    parser.add_argument('--model_dir', default='./src/checkpoints', type=str, help='directory to saved checkpoints')
    parser.add_argument('--output_dir', default='./src/outputs', type=str, help='directory to store outputs')

    # Experiment
    parser.add_argument('--start', default=0, type=int, help='start index of test images (after shuffling)')
    parser.add_argument('--num_images', default=1000, type=int, help='number of images to experiment')
    parser.add_argument('--batch_size', default=50, type=int, help='number of images per batch')
    parser.add_argument('--save', action='store_true', help='save original and safe images')
    parser.add_argument('--eval', action='store_true', help='evaluate original and safe images')

    # Dataset
    parser.add_argument('--data_type', default='CIFAR10', choices=['CIFAR10', 'ImageNet'], type=str, help='datasets')

    # Model
    parser.add_argument('--model_type', default='WideResNet', type=str, help='model architecture')
    parser.add_argument('--train_type', default='nat', type=str, help='model training method')
    
    # Safe spot
    parser.add_argument('--defense_type', default='SafeSpotLinf', choices=['SafeSpotLinf', 'SafeSpotL2', 'SafeSpotL2Rand', 'PGDAttackLinf', 'PGDAttackL2'], type=str)
    parser.add_argument('--delta', default=0.031, type=float, help='8/255=0.031, 4/255=0.0155')
    parser.add_argument('--num_iters', default=100, type=int, help='MAX_ITER. number of iterations for safe spot update')
    parser.add_argument('--lr', default=0.1, type=float)
    parser.add_argument('--hessian', action='store_true', help='use second-order gradients')

    # Attack (safe image update)
    parser.add_argument('--attack_type', default='PGDAttackLinf', choices=['PGDAttackLinf', 'PGDAttackL2', 'PGDAttackL2Rand'], type=str)
    parser.add_argument('--epsilon', default=0.031, type=float, help='8/255=0.031, 4/255=0.0155')
    parser.add_argument('--step_size', default=0.00775, type=float, help='2/255=0.00775, 1/255=0.003875')
    parser.add_argument('--num_steps', default=20, type=int)
    parser.add_argument('--random_starts', default=1, type=int, help='0 if no random start')
    parser.add_argument('--num_samples', default=5, type=int, help='number of sample when attacking smoothed classifiers')

    # Attack (safe image evaluation)
    parser.add_argument('--attack_type_test', default='PGDAttackLinf', choices=['PGDAttackLinf', 'PGDAttackL2', 'PGDAttackL2Rand'], type=str)
    parser.add_argument('--epsilon_test', default=0.031, type=float, help='8/255=0.031, 4/255=0.0155')
    parser.add_argument('--step_size_test', default=0.00775, type=float, help='2/255=0.00775, 1/255=0.003875')
    parser.add_argument('--num_steps_test', default=20, type=int)
    parser.add_argument('--random_starts_test', default=1, type=int, help='0 if no random start')
    parser.add_argument('--num_samples_test', default=50, type=int, help='number of sample when attacking smoothed classifiers')

    
    # Randomized smoothing
    parser.add_argument('--scale', default=0.05, type=float)
    parser.add_argument('--eval_rand', action='store_true', help='evaluate original and safe images with smoothed classifiers')

    args = parser.parse_args()

    # check argument conditions
    assert args.num_images % args.batch_size == 0 and args.start % args.batch_size == 0

    return args


def print_args(args):
    for key, val in vars(args).items():
        print('{:<20} : {}'.format(key, val))
