import argparse

parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
# Training settings
parser.add_argument('-b', '--batch-size', type=int, default=64, metavar='N',
                    help='input batch size for training (default: 64)')
parser.add_argument('--epochs', type=int, default=5, metavar='N',
                    help='number of epochs to train (default: 14)')
parser.add_argument('--lr', type=float, default=1.0, metavar='LR',
                    help='learning rate (default: 1.0)')
parser.add_argument('--min-lr', type=float, default=1e-6,
                    help='minimal learning rate (default: 1.0)')
parser.add_argument('--weight-decay', type=float, default=1e-3, metavar='WD',
                    help='weight decay (default: 1e-3)')
parser.add_argument('--momentum', type=float, default=0.9, metavar='MT',
                    help='momentum (default: 0.9)')
parser.add_argument('--num-workers', type=int, default=8, metavar='NW',
                    help='number of workers (default: 8)')
parser.add_argument('--gamma', type=float, default=0.7, metavar='M',
                    help='Learning rate step gamma (default: 0.7)')
parser.add_argument('--step-size', type=float, default=10,
                    help='Step size to update learning rate')
parser.add_argument('--no-cuda', action='store_true', default=False,
                    help='disables CUDA training')
parser.add_argument('--dry-run', action='store_true', default=False,
                    help='quickly check a single pass')
parser.add_argument('--adam', action='store_true', default=False,
                    help='using adam optim')
parser.add_argument('--adamW', action='store_true', default=False,
                    help='using adamW optim')
parser.add_argument('--sgd', action='store_true', default=False,
                    help='using SGD optim')
parser.add_argument('--cos', action='store_true', default=False,
                    help='using cosine scheduler')
parser.add_argument('--step', action='store_true', default=False,
                    help='using step scheduler')
parser.add_argument('--step_scheduler-size', type=int, default=1,
                    help='step size for step scheduler (default: 1)')
parser.add_argument('--multi-step', action='store_true', default=False,
                    help='using multi-step scheduler')
parser.add_argument('--seed', type=int, default=1, metavar='S',
                    help='random seed (default: 1)')
parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
                    help='input batch size for testing (default: 1000)')
parser.add_argument('--model-type', type=str, default='conv',
                    help='which kind of model to use')
parser.add_argument('--cudnn', action='store_true', default=False,
                    help='use cudnn benchmark to speed up the training')
parser.add_argument('--dev', action='store_true', default=False,
                    help='debug mode')
parser.add_argument('--log-img', action='store_true', default=False,
                    help='log images')
parser.add_argument('--amp', action='store_true', default=False,
                    help='use AMP training')
parser.add_argument('--scaler', action='store_true', default=False,
                    help='use gradient scaler training')
parser.add_argument('--bf16', action='store_true', default=False,
                    help='use bfloat16 for training')

# Log settings
parser.add_argument('--base-log-dir', type=str, default='./logs',
                    help='log output directory')
parser.add_argument('--log', action='store_true', default=False,
                    help='log output and train images')
parser.add_argument('--log-interval', type=int, default=200, metavar='N',
                    help='how many batches to wait before logging training status')
parser.add_argument('--save-model', action='store_true', default=False,
                    help='For saving the current Model')
parser.add_argument('--log-grad-norm', action='store_true', default=False,
                    help='log model gradient norm')
parser.add_argument('--exp', type=str, default='task',
                    help='name of the experiments')
parser.add_argument('--save-rec', action='store_true', default=False,
                    help='For saving the current Model recurrently')
parser.add_argument('--save-best', action='store_true', default=False,
                    help='For saving the current best Model')
parser.add_argument('--save-interval', type=int, default=5,
                    help='how many epochs to wait before save the model')
parser.add_argument('--load-model', type=str, default='',
                    help='load pre-trained model')
parser.add_argument('--resume', type=str, default='',
                    help='resume from previous training')
parser.add_argument('--cur-ep', type=int, default=None,
                    help='epoch of previous training')


# Augmentation
parser.add_argument('--rotate', action='store_true', default=False,
                    help='rotate the input')
parser.add_argument('--train-rotate', action='store_true', default=False,
                    help='rotate the input during training')
parser.add_argument('--img-size', type=int, default=128,
                    help='target image size of the input')
parser.add_argument('--degree', type=int, default=180,
                    help='degree to rotate during test time')
parser.add_argument('--train-degree', type=int, default=30,
                    help='degree to rotate during train time')
parser.add_argument('--fix-rotate', action='store_true', default=False,
                    help='rotate the input with fixed angle')
parser.add_argument('--moco-aug', action='store_true', default=False,
                    help='using MoCo style transforms')
parser.add_argument('--translation', action='store_true', default=False,
                    help='add random translation')
parser.add_argument('--test-translation', action='store_true', default=False,
                    help='add random translation')
parser.add_argument('--translate-ratio', type=float, default=0.2,
                    help='ratio of random translation')
parser.add_argument('--test-translate-ratio', type=float, default=None,
                    help='ratio of random translation at test time')
parser.add_argument('--vflip', action='store_true', default=False,
                    help='add vertical flipping')
parser.add_argument('--hflip', action='store_true', default=False,
                    help='add horizontal flipping')
parser.add_argument('--dflip', action='store_true', default=False,
                    help='add 3d flipping')
parser.add_argument('--eval-flip', action='store_true', default=False,
                    help='evaluate the model with flipping')
parser.add_argument('--eval-interval', type=int, default=1000,
                    help='interval of specific evaluation')
parser.add_argument('--pad-size-h', type=int, default=None,
                    help='height to pad the image')
parser.add_argument('--pad-size-w', type=int, default=None,
                    help='weight to pad the image')
parser.add_argument('--train-flip-p', type=float, default=0.5,
                    help='probability of apply flipping during training')
parser.add_argument('--affine-prob', type=float, default=0.2,
                    help='probability of apply affine transform during training')
parser.add_argument('--scale-ratio', type=float, default=0.1,
                    help='ratio to scale the image')
parser.add_argument('--shear-ratio', type=float, default=0.1,
                    help='ratio to shear the image')
parser.add_argument('--expand', action='store_true', default=False,
                    help='expand the image during rotation')
parser.add_argument('--padding', type=str, default='zeros',
                    help='padding method for rotation: zeros | reflect')
parser.add_argument('--aug-axis', type=str, default='x',
                    help='axis to apply augmentation: x | y | z')
parser.add_argument('--interpolation', type=str, default='nearest',
                    help='interpolation method for rotation: nearest | bilinear | bicubic')


# SRE conv settings
parser.add_argument('--sre-shape', type=str, default='o',
                    help='shape of the RI conv kernel')
parser.add_argument('--sre-conv-size', type=int, default=3,
                    help='size of the RI convolution kernel')
parser.add_argument('--sre-conv-k', nargs='+', default=None,
                    help='# of layers of the RI convolution kernel')
parser.add_argument('--res-inplanes', type=int, default=64,
                    help='use different convolution channel size')
parser.add_argument('--res-keep-conv1', action='store_true', default=False,
                    help='Not to down sample for resnet Conv1')

# Model settings
parser.add_argument('--dropout-p', type=float, default=-1,
                    help='probability of dropout layers')
parser.add_argument('--large-conv', action='store_true', default=False,
                    help='Use 7x7 convolution layers')
parser.add_argument('--maxpool', action='store_true', default=False,
                    help='use max pooling')




# MedMNIST args
parser.add_argument('--sre-conv-size-list', nargs='+', default=None,
                    help='Size of each MNIST ri conv kernel')
parser.add_argument('--eval-rot', action='store_true', default=False,
                    help='Evaluate trained model on different rotation degrees')
parser.add_argument('--med-mnist', type=str, default='',
                    help='indicate the medMNIST dataset to be used')
parser.add_argument('--no-aug', action='store_true', default=False,
                    help='Train w/o augmentation')
parser.add_argument('--multi-label', action='store_true', default=False,
                    help='Train with multi-label loss')
parser.add_argument('--medmnist-size', type=int, default=28,
                    help='Size of MedMNIST image')


def get_opt():
    args = parser.parse_args()
    if args.sre_conv_size_list != None:
        args.sre_conv_size_list = [int(x) for x in args.sre_conv_size_list]
    if isinstance(args.sre_conv_k, list):
        args.sre_conv_k = [int(x) for x in args.sre_conv_k]
    return args