"""
StarGAN v2
Copyright (c) 2020-present NAVER Corp.

This work is licensed under the Creative Commons Attribution-NonCommercial
4.0 International License. To view a copy of this license, visit
http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to
Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
"""

import os
import random
from os.path import join as ospj
import time
import datetime
from munch import Munch

import torch
import torch.nn as nn
import torch.nn.functional as F

from core.model import build_model
from core.stylegan_model import StyledGenerator, Discriminator
from core.checkpoint import CheckpointIO
from core.data_loader import InputFetcher
import core.utils as utils
from metrics.eval import calculate_metrics


fea_w={128:1., 64:1., 32:1., 16:1., 8:1., 4:1.}

class Solver(nn.Module):
    def __init__(self, args):
        super().__init__()
        self.args = args
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        self.nets, self.nets_ema = build_model(args)




        # below setattrs are to make networks be children of Solver, e.g., for self.to(self.device)
        for name, module in self.nets.items():
            utils.print_network(module, name)
            setattr(self, name, module)
        for name, module in self.nets_ema.items():
            setattr(self, name + '_ema', module)

        if args.mode == 'train':
            self.optims = Munch()
            for net in self.nets.keys():
                if net == 'fan':
                    continue
                self.optims[net] = torch.optim.Adam(
                    params=self.nets[net].parameters(),
                    lr=args.f_lr if net == 'mapping_network' else args.lr,
                    betas=[args.beta1, args.beta2],
                    weight_decay=args.weight_decay)

            self.ckptios = [
                CheckpointIO(ospj(args.checkpoint_dir, '{:06d}_nets.ckpt'), **self.nets),
                CheckpointIO(ospj(args.checkpoint_dir, '{:06d}_nets_ema.ckpt'), **self.nets_ema),
                CheckpointIO(ospj(args.checkpoint_dir, '{:06d}_optims.ckpt'), **self.optims)]
        else:
            self.ckptios = [CheckpointIO(ospj(args.checkpoint_dir, '{:06d}_nets_ema.ckpt'), **self.nets_ema)]

        self.to(self.device)
        for name, network in self.named_children():
            # Do not initialize the FAN parameters
            if ('ema' not in name) and ('fan' not in name):
                print('Initializing %s...' % name)
                network.apply(utils.he_init)
        # styleGAN
        StyleGAN=True
        if StyleGAN:
            self.style_generator = nn.DataParallel(StyledGenerator(512)).cuda()
            self.style_generator.eval()
            # no_from_rgb_activate is False
            self.style_discriminator = nn.DataParallel(Discriminator(from_rgb_activate=True)).cuda()
            self.style_discriminator.eval()
            self.style_g_running = StyledGenerator(512).cuda()
            self.style_g_running.train(False)
            self.style_g_running.eval()
            stylegan_ckpt = torch.load(args.stylegan_ckpt)
            self.style_generator.module.load_state_dict(stylegan_ckpt['generator'])
            self.style_discriminator.module.load_state_dict(stylegan_ckpt['discriminator'])
            self.style_g_running.load_state_dict(stylegan_ckpt['g_running'])

    def _save_checkpoint(self, step):
        for ckptio in self.ckptios:
            ckptio.save(step)

    def _load_checkpoint(self, step):
        for ckptio in self.ckptios:
            ckptio.load(step)

    def _reset_grad(self):
        for optim in self.optims.values():
            optim.zero_grad()

    def train(self, loaders):
        args = self.args
        nets = self.nets
        nets_ema = self.nets_ema
        optims = self.optims

        # fetch random validation images for debugging
        fetcher = InputFetcher(loaders.src, loaders.ref, args.latent_dim, 'train')
        fetcher_val = InputFetcher(loaders.val, None, args.latent_dim, 'val')
        inputs_val = next(fetcher_val)

        # resume training if necessary

        if args.resume_iter > 0:
            self._load_checkpoint(args.resume_iter)

        # remember the initial value of ds weight
        initial_lambda_ds = args.lambda_ds

        print('Start training...')
        start_time = time.time()
        for i in range(args.resume_iter, args.total_iters):
            # fetch images and labels
            inputs = next(fetcher)
            x_real, y_org = inputs.x_src, inputs.y_src
            x_ref, x_ref2, y_trg = inputs.x_ref, inputs.x_ref2, inputs.y_ref
            z_trg, z_trg2 = inputs.z_trg, inputs.z_trg2

            masks = nets.fan.get_heatmap(x_real) if args.w_hpf > 0 else None

            # stylegan feature and output


            mixing=True
            y_org = y_org*0
            y_trg = y_trg*0
            b_size = x_real.size(0)
            if mixing and random.random() < 1.1: # mixing is True, since we need reference, thus we always use this one
                gen_in11, gen_in12, gen_in21, gen_in22 = torch.randn( 4, b_size, 512, device='cuda').chunk(4, 0)
                gen_in1 = [gen_in11.squeeze(0), gen_in12.squeeze(0)]
                gen_in2 = [gen_in21.squeeze(0), gen_in22.squeeze(0)]
            else:
                gen_in1, gen_in2 = torch.randn(2, b_size, 512, device='cuda').chunk( 2, 0)
                gen_in1 = gen_in1.squeeze(0)
                gen_in2 = gen_in2.squeeze(0)

            Fake_img_1_1, G_feature_1_1 = self.style_g_running(gen_in1)
            _, D_feature_1_1 = self.style_discriminator(Fake_img_1_1, step=6, alpha=1, E1_output_feat=True)
            Fake_img_1_2, G_feature_1_2 = self.style_g_running(gen_in1, reference=True)
            _, D_feature_1_2 = self.style_discriminator(Fake_img_1_2, step=6, alpha=1, E1_output_feat=True)
            mix_begin = random.randint(2,6)# in starganv2, reference image only provide the style information, thus here we have to keep the gen_in1 offer the pose information (the bottom style code)
            mix_end = random.randint(mix_begin,6)
            Fake_img_1_mix_1_2, G_feature_1_mix_1_2 = self.style_g_running(gen_in1, mixing_range=(mix_begin, mix_end))
            _, D_feature_1_mix_1_2 = self.style_discriminator(Fake_img_1_mix_1_2, step=6, alpha=1, E1_output_feat=True)
            set1 = (gen_in1, Fake_img_1_1, G_feature_1_1, D_feature_1_1, Fake_img_1_2, G_feature_1_2, D_feature_1_2, Fake_img_1_mix_1_2, G_feature_1_mix_1_2, D_feature_1_mix_1_2)

            Fake_img_2_1, G_feature_2_1 = self.style_g_running(gen_in2)
            _, D_feature_2_1 = self.style_discriminator(Fake_img_2_1, step=6, alpha=1, E1_output_feat=True)
            Fake_img_2_2, G_feature_2_2 = self.style_g_running(gen_in2, reference=True)
            _, D_feature_2_2 = self.style_discriminator(Fake_img_2_2, step=6, alpha=1, E1_output_feat=True)
            mix_begin = random.randint(2,6)# in starganv2, reference image only provide the style information, thus here we have to keep the gen_in1 offer the pose information (the bottom style code)
            mix_end = random.randint(mix_begin,6)
            Fake_img_2_mix_1_2, G_feature_2_mix_1_2 = self.style_g_running(gen_in2, mixing_range=(mix_begin, mix_end))
            _, D_feature_2_mix_1_2 = self.style_discriminator(Fake_img_2_mix_1_2, step=6, alpha=1, E1_output_feat=True)

            set2 = (gen_in2, Fake_img_2_1, G_feature_2_1, D_feature_2_1, Fake_img_2_2, G_feature_2_2, D_feature_2_2, Fake_img_2_mix_1_2, G_feature_2_mix_1_2, D_feature_2_mix_1_2)

            # train the discriminator
            d_loss, d_losses_latent = compute_d_loss(
                nets, args, x_real, y_org, y_trg, z_trg=z_trg, masks=masks, set1=set1, set2=set2)
            self._reset_grad()
            d_loss.backward(retain_graph=True)
            optims.discriminator.step()

            d_loss, d_losses_ref = compute_d_loss(
                nets, args, x_real, y_org, y_trg, x_ref=x_ref, masks=masks, set1=set1, set2=set2)
            self._reset_grad()
            d_loss.backward(retain_graph=True)
            optims.discriminator.step()

            # train the generator
            g_loss, g_losses_latent = compute_g_loss(
                nets, args, x_real, y_org, y_trg, z_trgs=[z_trg, z_trg2], masks=masks, set1=set1, set2=set2)
            self._reset_grad()
            g_loss.backward(retain_graph=True)
            optims.generator.step()
            optims.mapping_network.step()
            optims.style_encoder.step()

            g_loss, g_losses_ref = compute_g_loss(
                nets, args, x_real, y_org, y_trg, x_refs=[x_ref, x_ref2], masks=masks, set1=set1, set2=set2)
            self._reset_grad()
            g_loss.backward()
            optims.generator.step()

            # compute moving average of network parameters
            moving_average(nets.generator, nets_ema.generator, beta=0.)
            moving_average(nets.mapping_network, nets_ema.mapping_network, beta=0.)
            moving_average(nets.style_encoder, nets_ema.style_encoder, beta=0.)

            # decay weight for diversity sensitive loss
            if args.lambda_ds > 0:
                args.lambda_ds -= (initial_lambda_ds / args.ds_iter)

            # print out log info
            if (i+1) % args.print_every == 0:
                elapsed = time.time() - start_time
                elapsed = str(datetime.timedelta(seconds=elapsed))[:-7]
                log = "Elapsed time [%s], Iteration [%i/%i], " % (elapsed, i+1, args.total_iters)
                all_losses = dict()
                for loss, prefix in zip([d_losses_latent, d_losses_ref, g_losses_latent, g_losses_ref],
                                        ['D/latent_', 'D/ref_', 'G/latent_', 'G/ref_']):
                    for key, value in loss.items():
                        all_losses[prefix + key] = value
                all_losses['G/lambda_ds'] = args.lambda_ds
                log += ' '.join(['%s: [%.4f]' % (key, value) for key, value in all_losses.items()])
                print(log)

            # generate images for debugging
            if (i+1) % args.sample_every == 0:
                os.makedirs(args.sample_dir, exist_ok=True)
                utils.debug_image(nets_ema, args, inputs=inputs_val, step=i+1, set1=set1)

            # save model checkpoints
            if (i+1) % args.save_every == 0:
                self._save_checkpoint(step=i+1)

            # compute FID and LPIPS if necessary
            if (i+1) % args.eval_every == 0:
                calculate_metrics(nets_ema, args, i+1, mode='latent')
                calculate_metrics(nets_ema, args, i+1, mode='reference')

    @torch.no_grad()
    def sample(self, loaders):
        args = self.args
        nets_ema = self.nets_ema
        os.makedirs(args.result_dir, exist_ok=True)
        self._load_checkpoint(args.resume_iter)

        src = next(InputFetcher(loaders.src, None, args.latent_dim, 'test'))
        ref = next(InputFetcher(loaders.ref, None, args.latent_dim, 'test'))

        fname = ospj(args.result_dir, 'reference.jpg')
        print('Working on {}...'.format(fname))
        utils.translate_using_reference(nets_ema, args, src.x, ref.x, ref.y, fname)

        fname = ospj(args.result_dir, 'video_ref.mp4')
        print('Working on {}...'.format(fname))
        utils.video_ref(nets_ema, args, src.x, ref.x, ref.y, fname)

    @torch.no_grad()
    def evaluate(self):
        args = self.args
        nets_ema = self.nets_ema
        resume_iter = args.resume_iter
        self._load_checkpoint(args.resume_iter)
        calculate_metrics(nets_ema, args, step=resume_iter, mode='latent')
        calculate_metrics(nets_ema, args, step=resume_iter, mode='reference')


def downsampling(x, num_split=32):
    out = torch.zeros_like(x)
    x = torch.split(x,num_split,dim=1)
    for i in range(len(x)):
        out[:,i] = (torch.mean(x[i],dim=1).clone()).detach()
    return out[:,:len(x)]

def l2_loss(stylegan_feat,stargan_feat,RESOLUTION={128, 64, 32, 16, 8, 4}):
    
    loss_M=0.
    for keys in RESOLUTION:
            loss_M += fea_w[keys] * torch.mean(torch.abs(stylegan_feat[keys]- stargan_feat[keys]))
    return loss_M

def compute_d_loss(nets, args, x_real, y_org, y_trg, z_trg=None, x_ref=None, masks=None, set1=None, set2=None):
    assert (z_trg is None) != (x_ref is None)

    gen_in1, Fake_img_1_1, G_feature_1_1, D_feature_1_1, Fake_img_1_2, G_feature_1_2, D_feature_1_2, Fake_img_1_mix_1_2, G_feature_1_mix_1_2, D_feature_1_mix_1_2 = set1  
    gen_in2, Fake_img_2_1, G_feature_2_1, D_feature_2_1, Fake_img_2_2, G_feature_2_2, D_feature_2_2, Fake_img_2_mix_1_2, G_feature_2_mix_1_2, D_feature_2_mix_1_2 = set2 
    # with real images
    out, stargan_D_feature = nets.discriminator(Fake_img_1_mix_1_2, y_org)

    loss_real = adv_loss(out, 1)
    loss_reg = r1_reg(out, Fake_img_1_mix_1_2)

    # with fake images
    with torch.no_grad():
        if z_trg is not None:
            #num_split=z_trg.size(1)
            #z_trg = gen_in1[1] 
            #z_trg = downsampling(z_trg, num_split=int(z_trg.size(1) / num_split))

            s_trg = nets.mapping_network(z_trg, y_trg)
        else:  # x_ref is not None
            s_trg, stargan_Sty_Enc_feature = nets.style_encoder(Fake_img_1_2, y_trg)
        x_fake, stargan_Enc_feature, stargan_Dec_feature = nets.generator(Fake_img_1_1, s_trg, masks=masks)
    out,_ = nets.discriminator(x_fake, y_trg)
    loss_fake = adv_loss(out, 0)

    # reconstuction loss
    l2_dis = l2_loss(D_feature_1_mix_1_2, stargan_D_feature)

    loss = loss_real + loss_fake + args.lambda_reg * loss_reg + 0.*(l2_dis)
    return loss, Munch(real=loss_real.item(),
                       fake=loss_fake.item(),
                       l2_dis=l2_dis.item(),
                       reg=loss_reg.item())


def compute_g_loss(nets, args, x_real, y_org, y_trg, z_trgs=None, x_refs=None, masks=None, set1=None, set2=None):
    assert (z_trgs is None) != (x_refs is None)
    if z_trgs is not None:
        z_trg, z_trg2 = z_trgs
    if x_refs is not None:
        x_ref, x_ref2 = x_refs

    gen_in1, Fake_img_1_1, G_feature_1_1, D_feature_1_1, Fake_img_1_2, G_feature_1_2, D_feature_1_2, Fake_img_1_mix_1_2, G_feature_1_mix_1_2, D_feature_1_mix_1_2 = set1  
    gen_in2, Fake_img_2_1, G_feature_2_1, D_feature_2_1, Fake_img_2_2, G_feature_2_2, D_feature_2_2, Fake_img_2_mix_1_2, G_feature_2_mix_1_2, D_feature_2_mix_1_2 = set2 

    # adversarial loss
    if z_trgs is not None:
       # num_split=z_trg.size(1)
       # z_trg = gen_in1[0] if type(gen_in1) is list else gen_in1
       # z_trg = downsampling(z_trg, num_split=int(z_trg.size(1) / num_split))
        s_trg = nets.mapping_network(z_trg, y_trg)
    else:
        s_trg, stargan_Sty_Enc_feature  = nets.style_encoder(Fake_img_1_2, y_trg)

    x_fake, stargan_Enc_feature, stargan_Dec_feature  = nets.generator(Fake_img_1_1, s_trg, masks=masks)
    out,_ = nets.discriminator(x_fake, y_trg)
    loss_adv = adv_loss(out, 1)


    # reconstuction loss
    l2_enc = l2_loss(D_feature_1_1, stargan_Enc_feature, RESOLUTION={128, 64, 32, 16}) 
    l2_dec = l2_loss(G_feature_1_mix_1_2, stargan_Dec_feature, RESOLUTION={128, 64, 32, 16}) 
    l2_img = torch.mean(torch.abs(x_fake - Fake_img_1_mix_1_2))

    if z_trgs is not None:
        l2_sty_enc = 0. 
    else:
        l2_sty_enc = l2_loss(D_feature_1_2, stargan_Sty_Enc_feature) 


    # style reconstruction loss
    s_pred, _ = nets.style_encoder(x_fake, y_trg)
    loss_sty = torch.mean(torch.abs(s_pred - s_trg))

    # diversity sensitive loss
    if z_trgs is not None:
       # num_split=z_trg2.size(1)
       # z_trg2 = gen_in2[0] if type(gen_in2) is list else gen_in2
       # z_trg2 = downsampling(z_trg2, num_split=int(z_trg2.size(1) / num_split))
        s_trg2 = nets.mapping_network(z_trg2, y_trg)
    else:
        s_trg2,_ = nets.style_encoder(Fake_img_2_2, y_trg)

    x_fake2,_,_ = nets.generator(Fake_img_1_1, s_trg2, masks=masks)
    x_fake2 = x_fake2.detach()
    loss_ds = torch.mean(torch.abs(x_fake - x_fake2))

    # cycle-consistency loss
    masks = nets.fan.get_heatmap(x_fake) if args.w_hpf > 0 else None
    s_org,_ = nets.style_encoder(Fake_img_1_1, y_org)
    x_rec,_,_ = nets.generator(x_fake, s_org, masks=masks)
    loss_cyc = torch.mean(torch.abs(x_rec - Fake_img_1_1))


    loss = loss_adv + args.lambda_sty * loss_sty \
        - args.lambda_ds * loss_ds + args.lambda_cyc * loss_cyc + 0.1*(l2_enc + l2_dec + l2_sty_enc + l2_img)
    return loss, Munch(adv=loss_adv.item(),
                       sty=loss_sty.item(),
                       ds=loss_ds.item(),
                       l2_enc=l2_enc.item(),
                       l2_dec=l2_dec.item(),
                       l2_sty_enc=l2_sty_enc.item() if z_trgs is None else 8888,
                       l2_img=l2_img.item(),
                       cyc=loss_cyc.item())


def moving_average(model, model_test, beta=0.999):
    for param, param_test in zip(model.parameters(), model_test.parameters()):
        param_test.data = torch.lerp(param.data, param_test.data, beta)


def adv_loss(logits, target):
    assert target in [1, 0]
    targets = torch.full_like(logits, fill_value=target)
    loss = F.binary_cross_entropy_with_logits(logits, targets)
    return loss


def r1_reg(d_out, x_in):
    # zero-centered gradient penalty for real images
    batch_size = x_in.size(0)
    grad_dout = torch.autograd.grad(
        outputs=d_out.sum(), inputs=x_in,
        create_graph=True, retain_graph=True, only_inputs=True
    )[0]
    grad_dout2 = grad_dout.pow(2)
    assert(grad_dout2.size() == x_in.size())
    reg = 0.5 * grad_dout2.view(batch_size, -1).sum(1).mean(0)
    return reg
