import argparse
import math
import torch

from gfpgan.archs.gfpganv1_clean_arch import GFPGANv1Clean


def modify_checkpoint(checkpoint_bilinear, checkpoint_clean):
    for ori_k, ori_v in checkpoint_bilinear.items():
        if 'stylegan_decoder' in ori_k:
            if 'style_mlp' in ori_k:  # style_mlp_layers
                lr_mul = 0.01
                prefix, name, idx, var = ori_k.split('.')
                idx = (int(idx) * 2) - 1
                crt_k = f'{prefix}.{name}.{idx}.{var}'
                if var == 'weight':
                    _, c_in = ori_v.size()
                    scale = (1 / math.sqrt(c_in)) * lr_mul
                    crt_v = ori_v * scale * 2**0.5
                else:
                    crt_v = ori_v * lr_mul * 2**0.5
                checkpoint_clean[crt_k] = crt_v
            elif 'modulation' in ori_k:  # modulation in StyleConv
                lr_mul = 1
                crt_k = ori_k
                var = ori_k.split('.')[-1]
                if var == 'weight':
                    _, c_in = ori_v.size()
                    scale = (1 / math.sqrt(c_in)) * lr_mul
                    crt_v = ori_v * scale
                else:
                    crt_v = ori_v * lr_mul
                checkpoint_clean[crt_k] = crt_v
            elif 'style_conv' in ori_k:
                # StyleConv in style_conv1 and style_convs
                if 'activate' in ori_k:  # FusedLeakyReLU
                    # eg. style_conv1.activate.bias
                    # eg. style_convs.13.activate.bias
                    split_rlt = ori_k.split('.')
                    if len(split_rlt) == 4:
                        prefix, name, _, var = split_rlt
                        crt_k = f'{prefix}.{name}.{var}'
                    elif len(split_rlt) == 5:
                        prefix, name, idx, _, var = split_rlt
                        crt_k = f'{prefix}.{name}.{idx}.{var}'
                    crt_v = ori_v * 2**0.5  # 2**0.5 used in FusedLeakyReLU
                    c = crt_v.size(0)
                    checkpoint_clean[crt_k] = crt_v.view(1, c, 1, 1)
                elif 'modulated_conv' in ori_k:
                    # eg. style_conv1.modulated_conv.weight
                    # eg. style_convs.13.modulated_conv.weight
                    _, c_out, c_in, k1, k2 = ori_v.size()
                    scale = 1 / math.sqrt(c_in * k1 * k2)
                    crt_k = ori_k
                    checkpoint_clean[crt_k] = ori_v * scale
                elif 'weight' in ori_k:
                    crt_k = ori_k
                    checkpoint_clean[crt_k] = ori_v * 2**0.5
            elif 'to_rgb' in ori_k:  # StyleConv in to_rgb1 and to_rgbs
                if 'modulated_conv' in ori_k:
                    # eg. to_rgb1.modulated_conv.weight
                    # eg. to_rgbs.5.modulated_conv.weight
                    _, c_out, c_in, k1, k2 = ori_v.size()
                    scale = 1 / math.sqrt(c_in * k1 * k2)
                    crt_k = ori_k
                    checkpoint_clean[crt_k] = ori_v * scale
                else:
                    crt_k = ori_k
                    checkpoint_clean[crt_k] = ori_v
            else:
                crt_k = ori_k
                checkpoint_clean[crt_k] = ori_v
            # end of 'stylegan_decoder'
        elif 'conv_body_first' in ori_k or 'final_conv' in ori_k:
            # key name
            name, _, var = ori_k.split('.')
            crt_k = f'{name}.{var}'
            # weight and bias
            if var == 'weight':
                c_out, c_in, k1, k2 = ori_v.size()
                scale = 1 / math.sqrt(c_in * k1 * k2)
                checkpoint_clean[crt_k] = ori_v * scale * 2**0.5
            else:
                checkpoint_clean[crt_k] = ori_v * 2**0.5
        elif 'conv_body' in ori_k:
            if 'conv_body_up' in ori_k:
                ori_k = ori_k.replace('conv2.weight', 'conv2.1.weight')
                ori_k = ori_k.replace('skip.weight', 'skip.1.weight')
            name1, idx1, name2, _, var = ori_k.split('.')
            crt_k = f'{name1}.{idx1}.{name2}.{var}'
            if name2 == 'skip':
                c_out, c_in, k1, k2 = ori_v.size()
                scale = 1 / math.sqrt(c_in * k1 * k2)
                checkpoint_clean[crt_k] = ori_v * scale / 2**0.5
            else:
                if var == 'weight':
                    c_out, c_in, k1, k2 = ori_v.size()
                    scale = 1 / math.sqrt(c_in * k1 * k2)
                    checkpoint_clean[crt_k] = ori_v * scale
                else:
                    checkpoint_clean[crt_k] = ori_v
                if 'conv1' in ori_k:
                    checkpoint_clean[crt_k] *= 2**0.5
        elif 'toRGB' in ori_k:
            crt_k = ori_k
            if 'weight' in ori_k:
                c_out, c_in, k1, k2 = ori_v.size()
                scale = 1 / math.sqrt(c_in * k1 * k2)
                checkpoint_clean[crt_k] = ori_v * scale
            else:
                checkpoint_clean[crt_k] = ori_v
        elif 'final_linear' in ori_k:
            crt_k = ori_k
            if 'weight' in ori_k:
                _, c_in = ori_v.size()
                scale = 1 / math.sqrt(c_in)
                checkpoint_clean[crt_k] = ori_v * scale
            else:
                checkpoint_clean[crt_k] = ori_v
        elif 'condition' in ori_k:
            crt_k = ori_k
            if '0.weight' in ori_k:
                c_out, c_in, k1, k2 = ori_v.size()
                scale = 1 / math.sqrt(c_in * k1 * k2)
                checkpoint_clean[crt_k] = ori_v * scale * 2**0.5
            elif '0.bias' in ori_k:
                checkpoint_clean[crt_k] = ori_v * 2**0.5
            elif '2.weight' in ori_k:
                c_out, c_in, k1, k2 = ori_v.size()
                scale = 1 / math.sqrt(c_in * k1 * k2)
                checkpoint_clean[crt_k] = ori_v * scale
            elif '2.bias' in ori_k:
                checkpoint_clean[crt_k] = ori_v

    return checkpoint_clean


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--ori_path', type=str, help='Path to the original model')
    parser.add_argument('--narrow', type=float, default=1)
    parser.add_argument('--channel_multiplier', type=float, default=2)
    parser.add_argument('--save_path', type=str)
    args = parser.parse_args()

    ori_ckpt = torch.load(args.ori_path)['params_ema']

    net = GFPGANv1Clean(
        512,
        num_style_feat=512,
        channel_multiplier=args.channel_multiplier,
        decoder_load_path=None,
        fix_decoder=False,
        # for stylegan decoder
        num_mlp=8,
        input_is_latent=True,
        different_w=True,
        narrow=args.narrow,
        sft_half=True)
    crt_ckpt = net.state_dict()

    crt_ckpt = modify_checkpoint(ori_ckpt, crt_ckpt)
    print(f'Save to {args.save_path}.')
    torch.save(dict(params_ema=crt_ckpt), args.save_path, _use_new_zipfile_serialization=False)