import argparse
import random

def args_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument('-m', default='cnncifar', type=str, help='Model')
    parser.add_argument('-d', default='cifar10', type=str, help='Dataset')

    parser.add_argument('-T', default=5, type=int, help="Number of time steps")
    parser.add_argument('-N', default=30, type=int, help="Number of clients")
    parser.add_argument('-M', default=60, type=int, help="Number of states")
    parser.add_argument('-S', default=10, type=int, help="Number of states per client")

    parser.add_argument('-R', default=100, type=int, help='Number of total training rounds')
    parser.add_argument('-K', default=5, type=int, help='Number of local steps')

    parser.add_argument('--set', default='partial', type=str, choices=['full', 'partial'], help='State setting')
    parser.add_argument('--partition', default='dir', type=str, choices=['dir', 'iid', 'exdir'], help='Data partition')
    parser.add_argument('--beta', default=100.0, type=float, nargs='*', help='The parameter `beta` of dirichlet distribution')
    parser.add_argument('--optim', default='sgd', type=str, choices=['sgd', 'adam'], help='Optimizer')
    parser.add_argument('--lr', default=0.01, type=float, help='Client/Local learning rate')
    parser.add_argument('--lr-decay', default=1.0, type=float, help='Learning rate decay')
    parser.add_argument('--momentum', default=0.0, type=float, help='Momentum of client optimizer')
    parser.add_argument('--weight-decay', default=1e-4, type=float, help='Weight decay of client optimizer')
    parser.add_argument('--global-lr', default=1.0, type=float, help='Server/Global learning rate')
    parser.add_argument('--batch-size', default=64, type=int, help='Mini-batch size')
    parser.add_argument('--seed', default=1234, type=int, help='Seed')
    parser.add_argument('--clip', default=0, type=int, help='Clip')
    parser.add_argument('--log', default='Log', type=str, help='Log, Log/Print')
    parser.add_argument('--device', default=0, type=int, help='Device')
    parser.add_argument('--save-model', default=0, type=int, help='Whether to save model')
    parser.add_argument('--capacity', default=500, type=int, help='Data capacity')

    parser.add_argument('--method', default="sfedpo", type=str, choices=['sfedpo', 'random'], help='Data selection method')
    parser.add_argument('--agg', default="sfedpo", type=str, choices=['sfedpo', 'avg'], help='Aggregation strategy')
    parser.add_argument('--distance', default="kl", type=str, choices=['l1', 'l2', 'kl', 'js'], help='Distance')
    parser.add_argument('--a1', default=0.1, type=float, help='a1')
    parser.add_argument('--b1', default=0.25, type=float, help='b1')
    parser.add_argument('--alpha', default=0.5, type=float, help='sampling budget')
    parser.add_argument('--a2', default=2, type=float, help='a2')
    parser.add_argument('--b2', default=0.0, type=float, help='b2')
    parser.add_argument('--num-classes', default=10, type=int, help='Number of classes')

    
    parser.add_argument('-a', default=1.0, type=float, help='a')
    parser.add_argument('-b', default=1.0, type=float, help='b')

    parser.add_argument('--dynamic', default="on", type=str, choices=['on', 'off'], help='Dynamic')

    parser.add_argument('--oracle', default="on", type=str, choices=['on', 'off'], help='Oracle')
    parser.add_argument('--epsilon', default=1.0, type=float, help='epsilon')
    args = parser.parse_args()
    
    return args
