from argparse import ArgumentParser

def parse_args():
    parser = ArgumentParser()
    parser.add_argument('--root', type=str, default='./data')
    parser.add_argument('--dataset', type=str, default='cifar100', choices=['cifar100', 'cifar10', 'imagenet'])
    parser.add_argument('--model', type=str, default='batch_pca', choices=['maha_ipca', 'batch_pca'])
    parser.add_argument('--task_type', type=str, default='standardCL_randomcls', choices=['standardCL_randomcls'], help='learning scenarios')
    parser.add_argument('--seed', type=int, default=0)
    parser.add_argument('--n_tasks', type=int, default=5)
    parser.add_argument('--validation', type=float, default=None, help='Propertion of dataset used e.g. if set 0.9, 90\% of training data is used for training and rest 10\% is used for validation')
    parser.add_argument('--optim_type', type=str, default='adam', choices=['adam', 'sgd'])
    parser.add_argument('--n_epochs', type=int, default=1)
    parser.add_argument('--loss_f', type=str, default='ce', choices=['ce', 'bce', 'nll'])
    parser.add_argument('--confusion', action='store_true')
    parser.add_argument('--tsne', action='store_true')
    parser.add_argument('--clip_init', action='store_true', help="Use CLIP encoder to initialize parameters. Note: results are not different without it. If used, must apply for baselines as well")
    parser.add_argument('--normalize', action='store_true', help="normalize parameters and features")
    parser.add_argument('--n_components', type=int, default=5)
    parser.add_argument('--folder', type=str, default=None, help='directory NAME. e.g. save under ./logs/NAME')
    parser.add_argument('--ff', type=float, default=1., help='diminishing weight on previous statistics')
    parser.add_argument('--dynamic', type=int, default=None, help='Set the max memory size. If set, use dynamic memory.')

    # Network
    parser.add_argument('--in_dim', type=int, default=512)
    parser.add_argument('--out_dim', type=int, default=1)

    # DataLoader
    parser.add_argument('--pin_memory', action='store_false')
    parser.add_argument('--num_workers', type=int, default=15)
    parser.add_argument('--lr', type=float, default=0.01)
    parser.add_argument('--batch_size', type=int, default=16)
    parser.add_argument('--minibatch_size', type=int, default=16, help="Batch size for pseudo-replay")
    parser.add_argument('--test_batch_size', type=int, default=512)
    
    args = parser.parse_args()
    return args
