"""
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, 45, 46, 59), 'direction_utils.load_G', 'load_G', ({(46, 52, 46, 58): 'config'}, {}), '(config)', False, 'from direction_utils import visualize_directions, load_G, get_direction_padding_fn, init_wandb, download_G\n'), ((53, 15, 53, 40), 'torch.cuda.device_count', 'torch.cuda.device_count', ({}, {}), '()', False, 'import torch\n'), ((57, 10, 57, 42), 'direction_utils.get_direction_padding_fn', 'get_direction_padding_fn', ({(57, 35, 57, 41): 'config'}, {}), '(config)', False, 'from direction_utils import visualize_directions, load_G, get_direction_padding_fn, init_wandb, download_G\n'), ((70, 12, 70, 59), 'torch.optim.Adam', 'torch.optim.Adam', (), '', False, 'import torch\n'), ((74, 13, 75, 68), 'utils.prepare_z_y', 'utils.prepare_z_y', (), '', False, 'import utils\n'), ((78, 23, 80, 63), 'utils.prepare_z_y', 'utils.prepare_z_y', (), '', False, 'import utils\n'), ((84, 25, 86, 65), 'utils.prepare_z_y', 'utils.prepare_z_y', (), '', False, 'import utils\n'), ((108, 13, 108, 76), 'torch.utils.tensorboard.SummaryWriter', 'SummaryWriter', ({(108, 27, 108, 75): "'%s/%s' % (config['logs_root'], experiment_name)"}, {}), "('%s/%s' % (config['logs_root'], experiment_name))", False, 'from torch.utils.tensorboard import SummaryWriter\n'), ((109, 19, 110, 81), 'train_fns.save_and_sample', 'train_fns.save_and_sample', ({(109, 45, 109, 55): 'G.module.G', (109, 57, 109, 61): 'None', (109, 63, 109, 73): 'G.module.G', (109, 75, 109, 77): 'z_', (109, 79, 109, 81): 'y_', (109, 83, 109, 90): 'fixed_z', (109, 92, 109, 99): 'fixed_y', (110, 45, 110, 55): 'state_dict', (110, 57, 110, 63): 'config', (110, 65, 110, 80): 'experiment_name'}, {}), '(G.module.G, None, G.module.G, z_, y_, fixed_z,\n    fixed_y, state_dict, config, experiment_name)', False, 'import train_fns\n'), ((169, 13, 169, 35), 'utils.prepare_parser', 'utils.prepare_parser', ({}, {}), '()', False, 'import utils\n'), ((33, 18, 33, 83), 'orojar.orojar', 'orojar', (), '', False, 'from orojar import orojar\n'), ((42, 8, 42, 89), 'direction_utils.init_wandb', 'init_wandb', ({(42, 19, 42, 25): 'config', (42, 27, 42, 52): "config['experiment_name']", (42, 54, 42, 76): "config['wandb_entity']", (42, 78, 42, 88): '"""imagenet"""'}, {}), "(config, config['experiment_name'], config['wandb_entity'],\n    'imagenet')", False, 'from direction_utils import visualize_directions, load_G, get_direction_padding_fn, init_wandb, download_G\n'), ((44, 8, 44, 20), '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, 21, 123, 68), '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, 12, 51, 41), 'sync_batchnorm.patch_replication_callback', 'patch_replication_callback', ({(51, 39, 51, 40): 'G'}, {}), '(G)', False, 'from sync_batchnorm import patch_replication_callback\n'), ((62, 31, 62, 88), 'torch.eye', 'torch.eye', (), '', False, 'import torch\n'), ((66, 8, 66, 40), 'torch.nn.init.kaiming_normal_', 'torch.nn.init.kaiming_normal_', ({(66, 38, 66, 39): 'A'}, {}), '(A)', False, 'import torch\n'), ((102, 27, 102, 53), 'torch.eye', 'torch.eye', ({(102, 37, 102, 52): "config['ndirs']"}, {}), "(config['ndirs'])", False, 'import torch\n'), ((129, 19, 129, 113), 'utils.progress', 'utils.progress', (), '', False, 'import utils\n'), ((131, 19, 131, 37), 'tqdm.tqdm', 'tqdm', ({(131, 24, 131, 36): 'dummy_loader'}, {}), '(dummy_loader)', False, 'from tqdm import tqdm\n'), ((140, 16, 140, 60), 'torch.zeros', 'torch.zeros', ({(140, 28, 140, 59): "(G_batch_size, config['ndirs'])"}, {}), "((G_batch_size, config['ndirs']))", False, 'import torch\n'), ((65, 31, 65, 90), 'torch.empty', 'torch.empty', (), '', False, 'import torch\n'), ((118, 20, 118, 40), 'layers.fast_gram_schmidt', 'fast_gram_schmidt', ({(118, 38, 118, 39): 'A'}, {}), '(A)', False, 'from layers import fast_gram_schmidt, norm\n'), ((159, 33, 160, 80), '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, 28, 147, 48), 'layers.fast_gram_schmidt', 'fast_gram_schmidt', ({(147, 46, 147, 47): 'A'}, {}), '(A)', False, 'from layers import fast_gram_schmidt, norm\n')]