import torch
import argparse

from iobm.container import cGAN, cGAN_train_configs

def parse_arguments():
    parser = argparse.ArgumentParser(description='cGAN Configuration and Training')

    parser.add_argument('--data', type=str, required=True, help='Directory name containing the data')
    parser.add_argument('--model', type=str, required=False, help='Pretrained to load. Leave blank to initialize model')
    parser.add_argument('--epochs', type=int, required=True, help='Number of epochs to train')
    parser.add_argument('--batch_size', type=int, default=64, help='Batch size for training data')

    return parser.parse_args()

# Necessary code
args = parse_arguments()
configs = cGAN_train_configs(args)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"\nFound {configs.n_classes} possible classes of data: {configs.data_name}")

if torch.cuda.is_available():
    print("Using GPU : ", torch.cuda.get_device_name(0))
else:
    print("No GPU available, using CPU.")

def run_cGAN_training() -> None:

    trainer = cGAN(
        device=device,
        data_name=configs.data_name,
        n_classes=configs.n_classes,
        input_model=configs.input_model,
        project_path=configs.project_path,
        latent_size=configs.latent_size,
        embedding_size=configs.embedding_size,
        batch_size=configs.batch_size,
        generator_lr=configs.generator_lr,
        discriminator_lr=configs.discriminator_lr,
        lambda_gp=configs.lambda_gp
    )
    trainer.train(num_epochs=configs.epochs)

if __name__ == "__main__":
    run_cGAN_training()