"""
Learns a matrix of Z-Space directions using a pre-trained BigGAN Generator.
Modified from train.py in the PyTorch BigGAN repo.
"""
import os
from tqdm import tqdm

import torch
import torch.nn as nn
import torch.optim

import utils
import train_fns
from sync_batchnorm import patch_replication_callback

from torch.utils.tensorboard import SummaryWriter

from orojar import orojar
from direction_utils import visualize_directions, load_G, get_direction_padding_fn, init_wandb, download_G
from layers import fast_gram_schmidt, norm


class DataParallelLoss(nn.Module):
    """
    This is simply a wrapper class to compute the OroJaR efficiently over several GPUs
    """

    def __init__(self, G):
        super(DataParallelLoss, self).__init__()
        self.G = G

    def forward(self, z, y, w, Q):
        penalty = orojar(self.G, z, c=y, w=w, G_z=None, Q=Q, multiple_layers=False)

        return penalty


# The main training file. Config is a dictionary specifying the configuration
# of this training run.
def run(config):
    if config['wandb_entity'] is not None:
        init_wandb(config, config['experiment_name'], config['wandb_entity'], 'imagenet')
    if config["G_path"] is None:  # Download a pre-trained G if necessary
        download_G()
        config["G_path"] = 'checkpoints/138k'
    G, state_dict, device, experiment_name = load_G(config)
    # If parallel, parallelize the GD module
    if config['parallel']:
        G = nn.DataParallel(DataParallelLoss(G))
        if config['cross_replica']:
            patch_replication_callback(G)

    num_gpus = torch.cuda.device_count()
    print(f'Using {num_gpus} GPUs')

    # If search_space != 'all', then we need to pad the z components that we are leaving alone:
    pad = get_direction_padding_fn(config)
    direction_size = config['dim_z'] if config['search_space'] == 'all' else config['ndirs']
    # A is our (ndirs, |z|) matrix of directions, where ndirs indicates the number of directions we want to learn
    if config['load_A'] == 'coords':
        print('Initializing with standard basis directions')
        A = torch.nn.Parameter(torch.eye(config['ndirs'], direction_size, device=device), requires_grad=True)
    elif config['load_A'] == 'random':
        print('Initializing with random directions')
        A = torch.nn.Parameter(torch.empty(config['ndirs'], direction_size, device=device), requires_grad=True)
        torch.nn.init.kaiming_normal_(A)
    else:
        raise NotImplementedError
    # We only learn A; G is left frozen during training:
    optim = torch.optim.Adam(params=[A], lr=config['A_lr'])

    # Allow for different batch sizes in G
    G_batch_size = max(config['G_batch_size'], config['batch_size'])
    z_, y_ = utils.prepare_z_y(G_batch_size, G.module.G.dim_z, config['n_classes'],
                               device=device, fp16=config['G_fp16'])

    # Prepare a fixed z & y to see individual sample evolution throghout training
    fixed_z, fixed_y = utils.prepare_z_y(G_batch_size, G.module.G.dim_z,
                                         config['n_classes'], device=device,
                                         fp16=config['G_fp16'])
    fixed_z.sample_()
    fixed_y.sample_()

    interp_z, interp_y = utils.prepare_z_y(config["n_samples"], G.module.G.dim_z,
                                           config['n_classes'], device=device,
                                           fp16=config['G_fp16'])
    interp_z.sample_()
    interp_y.sample_()

    if config['fix_class'] is not None:
        y_ = y_.new_full(y_.size(), config['fix_class'])
        fixed_y = fixed_y.new_full(fixed_y.size(), config['fix_class'])
        interp_y = interp_y.new_full(interp_y.size(), config['fix_class'])

    print('Beginning training at epoch %d...' % state_dict['epoch'])
    # Train for specified number of epochs, although we mostly track G iterations.
    iters_per_epoch = 1000
    dummy_loader = [None] * iters_per_epoch  # We don't need any real data

    path_size = config['path_size']
    # Simply stores a |z|-dimensional one-hot vector indicating each direction we are learning:
    direction_indicators = torch.eye(config['ndirs']).to(device)

    G.eval()

    G.module.optim = optim

    writer = SummaryWriter('%s/%s' % (config['logs_root'], experiment_name))
    sample_sheet = train_fns.save_and_sample(G.module.G, None, G.module.G, z_, y_, fixed_z, fixed_y,
                                             state_dict, config, experiment_name)

    writer.add_image('samples', sample_sheet, 0)

    interp_y_ = G.module.G.shared(interp_y)

    norm_fn = norm
    # Make directions orthonormal via Gram Schmidt followed a normalization:
    Q = pad(norm_fn(fast_gram_schmidt(A))) if not config["no_ortho"] else pad(A)

    if config["vis_during_training"]:
        print("Generating initial visualizations...")
        interp_vis = visualize_directions(G.module.G, interp_z, interp_y_, path_sizes=path_size, Q=Q,
                                          high_quality=False, npv=1)
        for w_ix in range(config['ndirs']):
            writer.add_video('G_ema/w%03d' % w_ix, interp_vis[w_ix], 0, fps=24)

    for epoch in range(state_dict['epoch'], config['num_epochs']):
        if config['pbar'] == 'mine':
            pbar = utils.progress(dummy_loader, displaytype='s1k' if config['use_multiepoch_sampler'] else 'eta')
        else:
            pbar = tqdm(dummy_loader)
        for i, _ in enumerate(pbar):
            state_dict['itr'] += 1
            z_.sample_()
            if config['fix_class'] is None:
                y_.sample_()
            y = G.module.G.shared(y_)

            # OroJaR taken w.r.t. w_sampled, NOT z:
            w = torch.zeros((G_batch_size, config['ndirs'])) # equal to the one-hot w
            penalty = G(z_, y, w=w, Q=Q.repeat(num_gpus, 1)).mean()

            optim.zero_grad()
            penalty.backward()
            optim.step()
            # re-orthogonalize A for visualizations and the next training iteration:
            Q = pad(norm_fn(fast_gram_schmidt(A))) if not config["no_ortho"] else pad(A)

            # Log metrics to TensorBoard/WandB:
            cur_training_iter = epoch * iters_per_epoch + i
            writer.add_scalar(f'Metrics/orojar', penalty.item(), cur_training_iter)
            writer.add_scalar('Metrics/direction_norm', A.pow(2).mean().pow(0.5).item(), cur_training_iter)

            # Save directions and log visuals:
            if not (state_dict['itr'] % config['save_every']):
                torch.save(A.cpu().detach(), '%s/%s/A_%06d.pt' %
                           (config['weights_root'], experiment_name, cur_training_iter))
                if config["vis_during_training"]:
                    interp_vis = visualize_directions(G.module.G, interp_z, interp_y_, path_sizes=path_size, Q=Q,
                                                      high_quality=False, npv=1)
                    for w_ix in range(config['ndirs']):
                        writer.add_video('G_ema/w%03d' % w_ix, interp_vis[w_ix], cur_training_iter, fps=24)

        state_dict['epoch'] += 1


def main():
    # parse command line and run
    parser = utils.prepare_parser()
    config = vars(parser.parse_args())
    print(config)
    run(config)


if __name__ == '__main__':
    main()

[(46, 'direction_utils.load_G', 'load_G', False, 'from direction_utils import visualize_directions, load_G, get_direction_padding_fn, init_wandb, download_G\n'), (53, 'torch.cuda.device_count', 'torch.cuda.device_count', False, 'import torch\n'), (57, 'direction_utils.get_direction_padding_fn', 'get_direction_padding_fn', False, 'from direction_utils import visualize_directions, load_G, get_direction_padding_fn, init_wandb, download_G\n'), (70, 'torch.optim.Adam', 'torch.optim.Adam', False, 'import torch\n'), (74, 'utils.prepare_z_y', 'utils.prepare_z_y', False, 'import utils\n'), (78, 'utils.prepare_z_y', 'utils.prepare_z_y', False, 'import utils\n'), (84, 'utils.prepare_z_y', 'utils.prepare_z_y', False, 'import utils\n'), (108, 'torch.utils.tensorboard.SummaryWriter', 'SummaryWriter', False, 'from torch.utils.tensorboard import SummaryWriter\n'), (109, 'train_fns.save_and_sample', 'train_fns.save_and_sample', False, 'import train_fns\n'), (169, 'utils.prepare_parser', 'utils.prepare_parser', False, 'import utils\n'), (33, 'orojar.orojar', 'orojar', False, 'from orojar import orojar\n'), (42, 'direction_utils.init_wandb', 'init_wandb', False, 'from direction_utils import visualize_directions, load_G, get_direction_padding_fn, init_wandb, download_G\n'), (44, 'direction_utils.download_G', 'download_G', False, 'from direction_utils import visualize_directions, load_G, get_direction_padding_fn, init_wandb, download_G\n'), (122, 'direction_utils.visualize_directions', 'visualize_directions', False, 'from direction_utils import visualize_directions, load_G, get_direction_padding_fn, init_wandb, download_G\n'), (51, 'sync_batchnorm.patch_replication_callback', 'patch_replication_callback', False, 'from sync_batchnorm import patch_replication_callback\n'), (62, 'torch.eye', 'torch.eye', False, 'import torch\n'), (66, 'torch.nn.init.kaiming_normal_', 'torch.nn.init.kaiming_normal_', False, 'import torch\n'), (102, 'torch.eye', 'torch.eye', False, 'import torch\n'), (129, 'utils.progress', 'utils.progress', False, 'import utils\n'), (131, 'tqdm.tqdm', 'tqdm', False, 'from tqdm import tqdm\n'), (140, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (65, 'torch.empty', 'torch.empty', False, 'import torch\n'), (118, 'layers.fast_gram_schmidt', 'fast_gram_schmidt', False, 'from layers import fast_gram_schmidt, norm\n'), (159, 'direction_utils.visualize_directions', 'visualize_directions', False, 'from direction_utils import visualize_directions, load_G, get_direction_padding_fn, init_wandb, download_G\n'), (147, 'layers.fast_gram_schmidt', 'fast_gram_schmidt', False, 'from layers import fast_gram_schmidt, norm\n')]