import argparse
import os
import sys
import pickle
import math

import torch
import numpy as np
from torchvision import utils

from model import Generator, Discriminator


def convert_modconv(vars, source_name, target_name, flip=False):
    weight = vars[source_name + '/weight'].value().eval()
    mod_weight = vars[source_name + '/mod_weight'].value().eval()
    mod_bias = vars[source_name + '/mod_bias'].value().eval()
    noise = vars[source_name + '/noise_strength'].value().eval()
    bias = vars[source_name + '/bias'].value().eval()

    dic = {
        'conv.weight': np.expand_dims(weight.transpose((3, 2, 0, 1)), 0),
        'conv.modulation.weight': mod_weight.transpose((1, 0)),
        'conv.modulation.bias': mod_bias + 1,
        'noise.weight': np.array([noise]),
        'activate.bias': bias,
    }

    dic_torch = {}

    for k, v in dic.items():
        dic_torch[target_name + '.' + k] = torch.from_numpy(v)

    if flip:
        dic_torch[target_name + '.conv.weight'] = torch.flip(
            dic_torch[target_name + '.conv.weight'], [3, 4]
        )

    return dic_torch


def convert_conv(vars, source_name, target_name, bias=True, start=0):
    weight = vars[source_name + '/weight'].value().eval()

    dic = {'weight': weight.transpose((3, 2, 0, 1))}

    if bias:
        dic['bias'] = vars[source_name + '/bias'].value().eval()

    dic_torch = {}

    dic_torch[target_name + f'.{start}.weight'] = torch.from_numpy(dic['weight'])

    if bias:
        dic_torch[target_name + f'.{start + 1}.bias'] = torch.from_numpy(dic['bias'])

    return dic_torch


def convert_torgb(vars, source_name, target_name):
    weight = vars[source_name + '/weight'].value().eval()
    mod_weight = vars[source_name + '/mod_weight'].value().eval()
    mod_bias = vars[source_name + '/mod_bias'].value().eval()
    bias = vars[source_name + '/bias'].value().eval()

    dic = {
        'conv.weight': np.expand_dims(weight.transpose((3, 2, 0, 1)), 0),
        'conv.modulation.weight': mod_weight.transpose((1, 0)),
        'conv.modulation.bias': mod_bias + 1,
        'bias': bias.reshape((1, 3, 1, 1)),
    }

    dic_torch = {}

    for k, v in dic.items():
        dic_torch[target_name + '.' + k] = torch.from_numpy(v)

    return dic_torch


def convert_dense(vars, source_name, target_name):
    weight = vars[source_name + '/weight'].value().eval()
    bias = vars[source_name + '/bias'].value().eval()

    dic = {'weight': weight.transpose((1, 0)), 'bias': bias}

    dic_torch = {}

    for k, v in dic.items():
        dic_torch[target_name + '.' + k] = torch.from_numpy(v)

    return dic_torch


def update(state_dict, new):
    for k, v in new.items():
        if k not in state_dict:
            raise KeyError(k + ' is not found')

        if v.shape != state_dict[k].shape:
            raise ValueError(f'Shape mismatch: {v.shape} vs {state_dict[k].shape}')

        state_dict[k] = v


def discriminator_fill_statedict(statedict, vars, size):
    log_size = int(math.log(size, 2))

    update(statedict, convert_conv(vars, f'{size}x{size}/FromRGB', 'convs.0'))

    conv_i = 1

    for i in range(log_size - 2, 0, -1):
        reso = 4 * 2 ** i
        update(
            statedict,
            convert_conv(vars, f'{reso}x{reso}/Conv0', f'convs.{conv_i}.conv1'),
        )
        update(
            statedict,
            convert_conv(
                vars, f'{reso}x{reso}/Conv1_down', f'convs.{conv_i}.conv2', start=1
            ),
        )
        update(
            statedict,
            convert_conv(
                vars, f'{reso}x{reso}/Skip', f'convs.{conv_i}.skip', start=1, bias=False
            ),
        )
        conv_i += 1

    update(statedict, convert_conv(vars, f'4x4/Conv', 'final_conv'))
    update(statedict, convert_dense(vars, f'4x4/Dense0', 'final_linear.0'))
    update(statedict, convert_dense(vars, f'Output', 'final_linear.1'))

    return statedict


def fill_statedict(state_dict, vars, size):
    log_size = int(math.log(size, 2))

    for i in range(8):
        update(state_dict, convert_dense(vars, f'G_mapping/Dense{i}', f'style.{i + 1}'))

    update(
        state_dict,
        {
            'input.input': torch.from_numpy(
                vars['G_synthesis/4x4/Const/const'].value().eval()
            )
        },
    )

    update(state_dict, convert_torgb(vars, 'G_synthesis/4x4/ToRGB', 'to_rgb1'))

    for i in range(log_size - 2):
        reso = 4 * 2 ** (i + 1)
        update(
            state_dict,
            convert_torgb(vars, f'G_synthesis/{reso}x{reso}/ToRGB', f'to_rgbs.{i}'),
        )

    update(state_dict, convert_modconv(vars, 'G_synthesis/4x4/Conv', 'conv1'))

    conv_i = 0

    for i in range(log_size - 2):
        reso = 4 * 2 ** (i + 1)
        update(
            state_dict,
            convert_modconv(
                vars,
                f'G_synthesis/{reso}x{reso}/Conv0_up',
                f'convs.{conv_i}',
                flip=True,
            ),
        )
        update(
            state_dict,
            convert_modconv(
                vars, f'G_synthesis/{reso}x{reso}/Conv1', f'convs.{conv_i + 1}'
            ),
        )
        conv_i += 2

    for i in range(0, (log_size - 2) * 2 + 1):
        update(
            state_dict,
            {
                f'noises.noise_{i}': torch.from_numpy(
                    vars[f'G_synthesis/noise{i}'].value().eval()
                )
            },
        )

    return state_dict


if __name__ == '__main__':
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print('Using PyTorch device', device)

    parser = argparse.ArgumentParser()
    parser.add_argument('--repo', type=str, required=True)
    parser.add_argument('--gen', action='store_true')
    parser.add_argument('--disc', action='store_true')
    parser.add_argument('--channel_multiplier', type=int, default=2)
    parser.add_argument('path', metavar='PATH')

    args = parser.parse_args()

    sys.path.append(args.repo)

    import dnnlib
    from dnnlib import tflib

    tflib.init_tf()

    with open(args.path, 'rb') as f:
        generator, discriminator, g_ema = pickle.load(f)

    size = g_ema.output_shape[2]

    g = Generator(size, 512, 8, channel_multiplier=args.channel_multiplier)
    state_dict = g.state_dict()
    state_dict = fill_statedict(state_dict, g_ema.vars, size)

    g.load_state_dict(state_dict)

    latent_avg = torch.from_numpy(g_ema.vars['dlatent_avg'].value().eval())

    ckpt = {'g_ema': state_dict, 'latent_avg': latent_avg}

    if args.gen:
        g_train = Generator(size, 512, 8, channel_multiplier=args.channel_multiplier)
        g_train_state = g_train.state_dict()
        g_train_state = fill_statedict(g_train_state, generator.vars, size)
        ckpt['g'] = g_train_state

    if args.disc:
        disc = Discriminator(size, channel_multiplier=args.channel_multiplier)
        d_state = disc.state_dict()
        d_state = discriminator_fill_statedict(d_state, discriminator.vars, size)
        ckpt['d'] = d_state

    name = os.path.splitext(os.path.basename(args.path))[0]
    outpath = os.path.join(os.getcwd(), f'{name}.pt')
    print('Saving', outpath)
    try:
        torch.save(ckpt, outpath, _use_new_zipfile_serialization=False)
    except TypeError:
        torch.save(ckpt, outpath)
    

    print('Generating TF-Torch comparison images')
    batch_size = {256: 8, 512: 4, 1024: 2}
    n_sample = batch_size.get(size, 4)

    g = g.to(device)

    z = np.random.RandomState(0).randn(n_sample, 512).astype('float32')

    with torch.no_grad():
        img_pt, _ = g(
            [torch.from_numpy(z).to(device)],
            truncation=0.5,
            truncation_latent=latent_avg.to(device),
        )

    img_tf = g_ema.run(z, None, randomize_noise=False)
    img_tf = torch.from_numpy(img_tf).to(device)

    img_diff = ((img_pt + 1) / 2).clamp(0.0, 1.0) - ((img_tf.to(device) + 1) / 2).clamp(
        0.0, 1.0
    )

    img_concat = torch.cat((img_tf, img_pt, img_diff), dim=0)
    utils.save_image(
        img_concat, name + '.png', nrow=n_sample, normalize=True, range=(-1, 1)
    )
    print('Done')