# Path: awesome_gans/config.py
# def parse_args():
#     parser = ArgumentParser(description='Awesome-GANs Arguments')
# 
#     # dataset
#     parser.add_argument('--width', type=int, default=32, help='width of image')
#     parser.add_argument('--height', type=int, default=32, help='height of image')
#     parser.add_argument('--n_channels', type=int, default=3, help='number of channel of image')
#     parser.add_argument('--root_path', type=str, default='./', help='root path')
#     parser.add_argument(
#         '--dataset',
#         type=str,
#         default='cifar10',
#         choices=['mnist', 'cifar10', 'cifar100', 'div2k'],
#         help='type of dataset',
#     )
#     parser.add_argument('--mnist_path', type=str, default='mnist')
#     parser.add_argument('--fashion_mnist_path', type=str, default='fashion-mnist')
#     parser.add_argument('--cifar10_path', type=str, default='cifar10')
#     parser.add_argument('--cifar100_path', type=str, default='cifar100')
#     parser.add_argument('--celeba_path', type=str, default='CelebA')
#     parser.add_argument('--celeba_hq_path', type=str, default='CelebA-HQ')
#     parser.add_argument('--div2k_path', type=str, default='DIV2K')
#     parser.add_argument('--pix2pix_path', type=str, default='pix2pix')
#     parser.add_argument('--use_crop', type=bool, default=False, help='use image center crop')
#     parser.add_argument('--model_path', type=str, default='model')
#     parser.add_argument('--output_path', type=str, default='outputs')
# 
#     # misc
#     parser.add_argument('--mode', default='train', type=str, choices=['train', 'inference'])
#     parser.add_argument('--n_samples', default=100, type=int, help='number of image samples to generate')
#     parser.add_argument('--device', default='cuda', type=str, help='type of device', choices=['cpu', 'cuda'])
#     parser.add_argument('--n_threads', default=8, type=int, help='number of threads')
#     parser.add_argument('--seed', default=42, type=int, help='seed for reproducibility')
#     parser.add_argument('--log_interval', default=1000, type=int, help='intervals to log')
#     parser.add_argument('--save_interval', default=1000, type=int, help='intervals to save the model(s)')
#     parser.add_argument('--verbose', type=bool, default=True)
# 
#     return parser


from awesome_gans.config import parse_args


def get_config():
    parser = parse_args()

    # Model
    parser.add_argument('--bs', default=64, type=int, help='batch size')
    parser.add_argument('--epochs', default=50, type=int, help='epochs to train')
    parser.add_argument('--global_steps', default=5e4, type=int, help='iterations to train')
    parser.add_argument('--n_feats', default=64, type=int, help='number of convolution filters')
    parser.add_argument('--d_lr', default=1e-4, type=float, help='learning rate of discriminator')
    parser.add_argument('--g_lr', default=1e-4, type=float, help='learning rate of generator')
    parser.add_argument(
        '--d_opt', default='rmsprop', type=str, choices=['adam', 'rmsprop', 'sgd'], help='disciminator optimizer'
    )
    parser.add_argument('--beta1', default=0.0, type=float)
    parser.add_argument('--beta2', default=0.99, type=float)
    parser.add_argument(
        '--g_opt', default='rmsprop', type=str, choices=['adam', 'rmsprop', 'sgd'], help='generator optimizer'
    )
    parser.add_argument('--d_loss', default='wgan', type=str)
    parser.add_argument('--g_loss', default='wgan', type=str)
    parser.add_argument('--grad_clip', default=1e-2, type=float)
    parser.add_argument(
        '--n_critics', default=5, type=int, help='number of times to train critic(discriminator) per 1-iter generator'
    )
    parser.add_argument('--z_dims', default=128, type=int, help='batch size')


# Path: awesome_gans/config.py
# def parse_args():
#     parser = ArgumentParser(description='Awesome-GANs Arguments')
# 
#     # dataset
#     parser.add_argument('--width', type=int, default=32, help='width of image')
#     parser.add_argument('--height', type=int, default=32, help='height of image')
#     parser.add_argument('--n_channels', type=int, default=3, help='number of channel of image')
#     parser.add_argument('--root_path', type=str, default='./', help='root path')
#     parser.add_argument(
#         '--dataset',
#         type=str,
#         default='cifar10',
#         choices=['mnist', 'cifar10', 'cifar100', 'div2k'],
#         help='type of dataset',
#     )
#     parser.add_argument('--mnist_path', type=str, default='mnist')
#     parser.add_argument('--fashion_mnist_path', type=str, default='fashion-mnist')
#     parser.add_argument('--cifar10_path', type=str, default='cifar10')
#     parser.add_argument('--cifar100_path', type=str, default='cifar100')
#     parser.add_argument('--celeba_path', type=str, default='CelebA')
#     parser.add_argument('--celeba_hq_path', type=str, default='CelebA-HQ')
#     parser.add_argument('--div2k_path', type=str, default='DIV2K')
#     parser.add_argument('--pix2pix_path', type=str, default='pix2pix')
#     parser.add_argument('--use_crop', type=bool, default=False, help='use image center crop')
#     parser.add_argument('--model_path', type=str, default='model')
#     parser.add_argument('--output_path', type=str, default='outputs')
# 
#     # misc
#     parser.add_argument('--mode', default='train', type=str, choices=['train', 'inference'])
#     parser.add_argument('--n_samples', default=100, type=int, help='number of image samples to generate')
#     parser.add_argument('--device', default='cuda', type=str, help='type of device', choices=['cpu', 'cuda'])
#     parser.add_argument('--n_threads', default=8, type=int, help='number of threads')
#     parser.add_argument('--seed', default=42, type=int, help='seed for reproducibility')
#     parser.add_argument('--log_interval', default=1000, type=int, help='intervals to log')
#     parser.add_argument('--save_interval', default=1000, type=int, help='intervals to save the model(s)')
#     parser.add_argument('--verbose', type=bool, default=True)
# 
#     return parser


# Path: awesome_gans/wgan/config.py
from awesome_gans.config import parse_args



def get_config():
    parser = parse_args()

    # Model
    parser.add_argument('--bs', default=64, type=int, help='batch size')
    parser.add_argument('--epochs', default=50, type=int, help='epochs to train')
    parser.add_argument('--global_steps', default=5e4, type=int, help='iterations to train')
    parser.add_argument('--n_feats', default=64, type=int, help='number of convolution filters')
    parser.add_argument('--d_lr', default=1e-4, type=float, help='learning rate of discriminator')
    parser.add_argument('--g_lr', default=1e-4, type=float, help='learning rate of generator')
    parser.add_argument(
        '--d_opt', default='rmsprop', type=str, choices=['adam', 'rmsprop', 'sgd'], help='disciminator optimizer'
    )
    parser.add_argument('--beta1', default=0.0, type=float)
    parser.add_argument('--beta2', default=0.99, type=float)
    parser.add_argument(
        '--g_opt', default='rmsprop', type=str, choices=['adam', 'rmsprop', 'sgd'], help='generator optimizer'
    )
    parser.add_argument('--d_loss', default='wgan', type=str)
    parser.add_argument('--g_loss', default='wgan', type=str)
    parser.add_argument('--grad_clip', default=1e-2, type=float)
    parser.add_argument(
        '--n_critics', default=5, type=int, help='number of times to train critic(discriminator) per 1-iter generator'
    )
    parser.add_argument('--z_dims', default=128, type=int, help='batch size')


    return parser.parse_args()