"""main.py"""

import argparse

import numpy as np

from solver import Solver
from utils import str2bool

import os



def main(args):

    net = Solver(args)

    if args.train:
        #os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
        #os.environ['CUDA_VISIBLE_DEVICES'] = "6"
        import torch
        torch.backends.cudnn.enabled = True
        torch.backends.cudnn.benchmark = True
        #torch.cuda.set_device(2)
        seed = args.seed
        torch.manual_seed(seed)
        torch.cuda.manual_seed(seed)
        np.random.seed(seed)
        net.train()
    else:
        # Instead of calling traverse here (as before), we call reconstruct to see how the images are reconstructed.
        # We also need to compute the KL Divergence and reconstruction loss. (See if we should use different flags)
        # for this reconstruction (which needs decoder net) and just for KLD (just encoder is enough).

        #net.traverse()
        os.environ['CUDA_VISIBLE_DEVICES'] = ""
        import torch
        torch.backends.cudnn.enabled = True
        torch.backends.cudnn.benchmark = True
        if args.is_audit == True:
            net.audit()
        else:
            net.test()



if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='toy Beta-VAE')

    parser.add_argument('--train', default=True, type=str2bool, help='train or traverse')
    parser.add_argument('--is_audit', default=False, type=str2bool, help='audit or not')
    parser.add_argument('--audit_classes', default=2, type=int, help='number of audit classes')
    parser.add_argument('--only_random', default=False, type=str2bool, help='Use model to just generate random images')
    parser.add_argument('--demo_test', default=False, type=str2bool, help='Save individual pictures or not for downstream API demographic test')
    parser.add_argument('--seed', default=1, type=int, help='random seed')
    parser.add_argument('--use_cuda', default=True, type=str2bool, help='enable cuda')
    parser.add_argument('--max_iter', default=1e6, type=float, help='maximum training iteration')
    parser.add_argument('--batch_size', default=64, type=int, help='batch size')

    parser.add_argument('--z_dim', default=10, type=int, help='dimension of the representation z')
    parser.add_argument('--beta', default=4, type=float, help='beta parameter for KL-term in original beta-VAE')
    parser.add_argument('--objective', default='H', type=str, help='beta-vae objective proposed in Higgins et al. or Burgess et al. H/B')
    parser.add_argument('--model', default='H', type=str, help='model proposed in Higgins et al. or Burgess et al. H/B')
    parser.add_argument('--gamma', default=1000, type=float, help='gamma parameter for KL-term in understanding beta-VAE')
    parser.add_argument('--C_max', default=25, type=float, help='capacity parameter(C) of bottleneck channel')
    parser.add_argument('--C_stop_iter', default=1e5, type=float, help='when to stop increasing the capacity')
    parser.add_argument('--lr', default=1e-4, type=float, help='learning rate')
    parser.add_argument('--beta1', default=0.9, type=float, help='Adam optimizer beta1')
    parser.add_argument('--beta2', default=0.999, type=float, help='Adam optimizer beta2')

    parser.add_argument('--dset_dir', default='data', type=str, help='dataset directory')
    parser.add_argument('--dataset', default='CelebA', type=str, help='dataset name')
    parser.add_argument('--image_size', default=64, type=int, help='image size. now only (64,64) is supported')
    parser.add_argument('--num_workers', default=2, type=int, help='dataloader num_workers')

    parser.add_argument('--viz_on', default=True, type=str2bool, help='enable visdom visualization')
    parser.add_argument('--viz_name', default='main', type=str, help='visdom env name')
    parser.add_argument('--viz_port', default=8097, type=str, help='visdom port number')
    parser.add_argument('--save_output', default=True, type=str2bool, help='save traverse images and gif')
    parser.add_argument('--output_dir', default='outputs', type=str, help='output directory')

    parser.add_argument('--gather_step', default=1000, type=int, help='numer of iterations after which data is gathered for visdom')
    parser.add_argument('--display_step', default=10000, type=int, help='number of iterations after which loss data is printed and visdom is updated')
    parser.add_argument('--save_step', default=10000, type=int, help='number of iterations after which a checkpoint is saved')

    parser.add_argument('--ckpt_dir', default='checkpoints', type=str, help='checkpoint directory')
    parser.add_argument('--ckpt_name', default='last', type=str, help='load previous checkpoint. insert checkpoint filename')

    parser.add_argument('--attr', default='male', type=str, help='enter the attribute of CelebA to test upon.')
    parser.add_argument('--do_modify', default=False, type=str2bool, help='flag to say if the training objective needs to call API and change capacity')

    # MORE NEW ONES
    parser.add_argument('--modify_step', default=500001, type=int, help='at which step to modify objective to add more capacity to selected dimensions')
    parser.add_argument('--modify_increase_step', default=20000, type=int, help='if we do incremental increase of capacity then for how many steps to increase it')
    parser.add_argument('--modify_type', default='d', type=str, help='choose between (i)ncremental or (d)irect increase of capacity')

    #Conditioning
    parser.add_argument('--condition', default=False, type=str2bool, help='flag to decide on Condition VAE or not')

    args = parser.parse_args()

    main(args)