import sys, argparse
from gantools import ganbreeder
from gantools import biggan
from gantools import latent_space
from gantools import image_utils

# create entrypoints for cli tools
def main():
    ## handle args
    parser = argparse.ArgumentParser(description='GAN tools')
    # load from ganbreeder
    parser.add_argument('-u', '--username', help='Ganbreeder account email address/username.')
    parser.add_argument('-p', '--password', help='Ganbreeder account password.')
    parser.add_argument('-k', '--keys', nargs='+', help='Ganbreeder keys.')
    parser.add_argument('-n', '--nframes', metavar='N', type=int, help='Total number of frames in the final animation.', default=10)
    parser.add_argument('-b', '--nbatch', metavar='N', type=int, help='Number of frames in each \'batch\' \
            (note: the truncation value can only change once per batch. Don\'t fuck with this unless you know \
            what it does.).', default=1)
    parser.add_argument('-o', '--output-dir', help='Directory path for output images.')
    parser.add_argument('--prefix', help='File prefix for output images.')
    args = parser.parse_args()
    # validate args
    if args.keys and not (args.username and args.password):
        parser.error('The --keys argument requires a --username and --password to login to ganbreeder')
        sys.exit(1)

    # get animation keyframes from ganbreeder
    print('Downloading keyframe info from ganbreeder...')
    keyframes = ganbreeder.get_info_batch(args.username, args.password, args.keys)

    # interpolate path through input space
    print('Interpolating path through input space...')
    z_seq, label_seq, truncation_seq = latent_space.sequence_keyframes(keyframes, args.nframes, args.nbatch)

    # sample the GAN
    print('Loading bigGAN...')
    gan = biggan.BigGAN()
    print('Sampling from bigGAN...')
    ims = gan.sample(z_seq, label_seq, truncation_seq, args.nbatch)

    # save images to file
    path = '' if args.output_dir == None else str(args.output_dir)
    prefix = '' if args.prefix == None else str(args.prefix)
    print('Saving image files: '+path + prefix)
    image_utils.save_images(ims, output_dir=output_dir, prefix=prefix)

    print('Done.')