import os
import glob
import time
import argparse

model_names = ['RANet']

arg_parser = argparse.ArgumentParser(description='RANet Image classification')

exp_group = arg_parser.add_argument_group('exp', 'experiment setting')
exp_group.add_argument('--save', default='save/default-{}'.format(time.time()),
                       type=str, metavar='SAVE',
                       help='path to the experiment logging directory'
                       '(default: save/debug)')
exp_group.add_argument('--resume', action='store_true', default=False,
                       help='path to latest checkpoint (default: False)')
exp_group.add_argument('--evalmode', default='anytime',
                       choices=['anytime', 'dynamic', 'both'],
                       help='which mode to evaluate')
exp_group.add_argument('--evaluate-from', default='', type=str, metavar='PATH',
                       help='path to saved checkpoint (default: none)')
exp_group.add_argument('--print-freq', '-p', default=176, type=int,
                       metavar='N', help='print frequency (default: 100)')
exp_group.add_argument('--seed', default=0, type=int,
                       help='random seed')
exp_group.add_argument('--gpu', default='0', type=str, help='GPU available.')

# dataset related
data_group = arg_parser.add_argument_group('data', 'dataset setting')
data_group.add_argument('--data', metavar='D', default='cifar10',
                        choices=['cifar10', 'cifar100', 'ImageNet','MiNi_ImageNet'],
                        help='data to work on')
data_group.add_argument('--data-root', metavar='DIR', default='/data/cx/data',
                        help='path to dataset (default: data)')
data_group.add_argument('--use_valid', action='store_true', default=False,
                        help='use validation set or not')
data_group.add_argument('-j', '--workers', default=4, type=int, metavar='N',
                        help='number of data loading workers (default: 4)')

# model arch related
arch_group = arg_parser.add_argument_group('arch', 'model architecture setting')
arch_group.add_argument('--arch', type=str, default='RANet')
arch_group.add_argument('--reduction', default=0.5, type=float,
                        metavar='C', help='compression ratio of DenseNet'
                        ' (1 means dot\'t use compression) (default: 0.5)')

# msdnet config
arch_group.add_argument('--nBlocks', type=int, default=2)
arch_group.add_argument('--nChannels', type=int, default=16)
arch_group.add_argument('--growthRate', type=int, default=6)
arch_group.add_argument('--grFactor', default='4-2-1', type=str)
arch_group.add_argument('--bnFactor', default='4-2-1', type=str)
arch_group.add_argument('--block-step', type=int, default=2)
arch_group.add_argument('--scale_list', default='1-2-3', type=str)
arch_group.add_argument('--compress-factor', default=0.25, type=float)
arch_group.add_argument('--step', type=int, default=4)
arch_group.add_argument('--stepmode', type=str, default='even', choices=['even', 'lg'])
arch_group.add_argument('--bnAfter', action='store_true', default=True)

#
# training related
optim_group = arg_parser.add_argument_group('optimization', 'optimization setting')
optim_group.add_argument('--epochs', default=300, type=int, metavar='N',
                         help='number of total epochs to run (default: 300)')
optim_group.add_argument('--start-epoch', default=0, type=int, metavar='N',
                         help='manual epoch number (useful on restarts)')
optim_group.add_argument('-b', '--batch-size', default=64, type=int,
                         metavar='N', help='mini-batch size (default: 64)')
optim_group.add_argument('--optimizer', default='sgd',
                         choices=['sgd', 'rmsprop', 'adam'], metavar='N',
                         help='optimizer (default=sgd)')
optim_group.add_argument('--lr', '--learning-rate', default=0.1, type=float,
                         metavar='LR',
                         help='initial learning rate (default: 0.1)')
optim_group.add_argument('--lr-type', default='multistep', type=str, metavar='T',
                        help='learning rate strategy (default: multistep)',
                        choices=['cosine', 'multistep'])
optim_group.add_argument('--decay-rate', default=0.1, type=float, metavar='N',
                         help='decay rate of learning rate (default: 0.1)')
optim_group.add_argument('--momentum', default=0.9, type=float, metavar='M',
                         help='momentum (default=0.9)')
optim_group.add_argument('--weight-decay', '--wd', default=1e-4, type=float,
                         metavar='W', help='weight decay (default: 1e-4)')
####
optim_group.add_argument('--T', '--Temperature', default=1.0, type=float, help='The temperature used for regularization training')
optim_group.add_argument('--sel_class', default=[0], type=float, help='Determine whether conduct CDM module on c-th classifier')
arch_group.add_argument('--CDM', action='store_true', default=True,help='If it set True, The classifier output will be sent to SoftPlus activation')
arch_group.add_argument('--ori_fusion', action='store_true', default=False,help='Original evidence fusion method')
arch_group.add_argument('--our_fusion', action='store_true', default=False,help='Our improved evidence fusion method')
arch_group.add_argument('--avge_fusion', action='store_true', default=False,help='Average evidence fusion method')
data_group.add_argument('--evaluate_from', metavar='DIR', default='./save_models/model_best.pth.tar',
                        help='model path')
data_group.add_argument('--dynamic_out_path', metavar='DIR', default='./dnynamic_out',
                        help='path to dataset (default: data)')
data_group.add_argument('--cuda_', default='cuda:0',help='Selecting cuda for training or testing')
arch_group.add_argument('--Eval_', action='store_true', default=False,help='Testing model or not')
data_group.add_argument('--save_dir', metavar='DIR', default='./',
                        help='path to save output files')
arch_group.add_argument('--train_trick', action='store_true', default=False,help='Average evidence fusion method')
optim_group.add_argument('--trick_T1',  default=0.5, type=float, help='The temperature used for regularization training')
optim_group.add_argument('--trick_T2', default=1.0, type=float, help='The temperature used for regularization training')
optim_group.add_argument('--lamda', default=0.1, type=float, help='The lamda used for regularization training')
args = arg_parser.parse_args()

args.grFactor = list(map(int, args.grFactor.split('-')))
args.bnFactor = list(map(int, args.bnFactor.split('-')))
args.scale_list = list(map(int, args.scale_list.split('-')))
args.nScales = len(args.grFactor)

if args.use_valid:
    args.splits = ['train', 'val', 'test']
else:
    args.splits = ['train', 'val']

if args.data == 'cifar10':
    args.num_classes = 10
elif args.data == 'cifar100':
    args.num_classes = 100
else:
    args.num_classes = 1000
