import argparse

import os
import torch

from imaginaire.config import Config
from imaginaire.utils.cudnn import init_cudnn
from imaginaire.utils.dataset import get_test_dataloader
from imaginaire.utils.distributed import init_dist
from imaginaire.utils.gpu_affinity import set_affinity
from imaginaire.utils.io import get_checkpoint as get_checkpoint
from imaginaire.utils.logging import init_logging
from imaginaire.utils.trainer import \
    (get_model_optimizer_and_scheduler, set_random_seed)
import imaginaire.config


def parse_args():
    parser = argparse.ArgumentParser(description='Training')
    parser.add_argument('--config', required=True,
                        help='Path to the training config file.')
    parser.add_argument('--checkpoint', default='',
                        help='Checkpoint path.')
    parser.add_argument('--output_dir', required=True,
                        help='Location to save the image outputs')
    parser.add_argument('--logdir',
                        help='Dir for saving logs and models.')
    parser.add_argument('--seed', type=int, default=0,
                        help='Random seed.')
    parser.add_argument('--debug', action='store_true')
    args = parser.parse_args()
    return args


def main():
    args = parse_args()
    cfg = Config(args.config)
    imaginaire.config.DEBUG = args.debug

    if not hasattr(cfg, 'inference_args'):
        cfg.inference_args = None

    # Create log directory for storing training results.
    cfg.date_uid, cfg.logdir = init_logging(args.config, args.logdir)

    # Initialize cudnn.
    init_cudnn(cfg.cudnn.deterministic, cfg.cudnn.benchmark)

    # Initialize data loaders and models.
    net_G = get_model_optimizer_and_scheduler(cfg, seed=args.seed, generator_only=True)

    if args.checkpoint == '':
        raise NotImplementedError("No checkpoint is provided for inference!")

    # Load checkpoint.
    # trainer.load_checkpoint(cfg, args.checkpoint)
    checkpoint = torch.load(args.checkpoint, map_location='cpu')
    net_G.load_state_dict(checkpoint['net_G'])

    # Do inference.
    net_G = net_G.module
    net_G.eval()
    for name, param in net_G.named_parameters():
        param.requires_grad = False
    torch.cuda.empty_cache()
    device = torch.device('cuda')
    rng_cuda = torch.Generator(device=device)
    rng_cuda = rng_cuda.manual_seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)

    world_dir = os.path.join(args.output_dir)
    os.makedirs(world_dir, exist_ok=True)
    print('[PCGGenerator] Generating BEV scene representation...')
    os.system('python terrain_generator.py --size {} --seed {} --outdir {}'.format(net_G.voxel.sample_size, args.seed, world_dir))
    net_G.voxel.next_world(device, world_dir, checkpoint)
    cam_mode = cfg.inference_args.camera_mode
    current_outdir = os.path.join(world_dir, 'camera_{:02d}'.format(cam_mode))
    os.makedirs(current_outdir, exist_ok=True)
    os.makedirs(current_outdir, exist_ok=True)
    z = torch.empty(1, net_G.style_dims, dtype=torch.float32, device=device)
    z.normal_(generator=rng_cuda)
    net_G.inference_givenstyle(z, current_outdir, **vars(cfg.inference_args))

if __name__ == "__main__":
    main()