# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
#
# 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.

"""Minimal script for generating an image using pre-trained StyleGAN generator."""

import os
import pickle
import numpy as np
import PIL.Image
import dnnlib
import dnnlib.tflib as tflib
import config
import tensorflow as tf

def main():
    # Initialize TensorFlow.
    tflib.init_tf()

    # Load pre-trained network.
    url = 'https://drive.google.com/uc?id=1MEGjdvVpUsu1jB4zrXZN7Y4kBBOzizDQ' # karras2019stylegan-ffhq-1024x1024.pkl
    with dnnlib.util.open_url(url, cache_dir=config.cache_dir) as f:
        _G, _D, Gs = pickle.load(f)
        # _G = Instantaneous snapshot of the generator. Mainly useful for resuming a previous training run.
        # _D = Instantaneous snapshot of the discriminator. Mainly useful for resuming a previous training run.
        # Gs = Long-term average of the generator. Yields higher-quality results than the instantaneous snapshot.

    print(type(Gs))
    # Print network details.
    Gs.print_layers()

    # Pick latent vector.
    rnd = np.random.RandomState(6)
    latents = rnd.randn(1, Gs.input_shape[1])
    print("input shape", Gs.input_shape)

    # Generate image.
    fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
    images = Gs.run(latents, None, truncation_psi=0.7, randomize_noise=False, output_transform=fmt)

    # Save image.
    os.makedirs(config.result_dir, exist_ok=True)
    png_filename = os.path.join(config.result_dir, 'example.png')
    PIL.Image.fromarray(images[0], 'RGB').save(png_filename)

    print("\n\ncache:",  Gs._run_cache)
    keys = list(Gs._run_cache.keys())

    in_expr, out_expr = Gs._run_cache[keys[0]]
    print(in_expr, out_expr)

    target = PIL.Image.open("results/base.png")
    target_expr = tf.constant(np.float32(target))

    learning_rate = .1
    opt = tf.train.RMSPropOptimizer(learning_rate=learning_rate)
    #train = opt.minimize(total_loss, var_list=[in_expr[0]])
    out_gpu = Gs.get_output_for(*in_expr, return_as_list=True)
    print("\n\n\n$$$$$$$$$$$$$$$$\n>", out_gpu)

    #grads,_ = tf.gradients(pixelloss, in_expr)
    #print("grads", grads)
    #exit()
    #writer = tf.summary.FileWriter("output", sess.graph)
    #rnd = np.random.RandomState(5)
    #latents2 = rnd.randn(1, Gs.input_shape[1])
    #dx2 = latents2-latents
    tf.keras.backend.set_image_data_format('channels_first')

    #pretrained_resnet = tf.keras.applications.MobileNet(
    #    input_tensor = tf.stack([tf.reshape(out_gpu,(3,1024,1024)),target_expr],0),
        #input_tensor = tf.reshape(target_expr,(1,3,1024,1024)),
    #    weights="imagenet",
    #    include_top=False,
        #input_shape=(1024,1024,3)
    #)

    #l1,l2 = tf.split(pretrained_resnet.output, 2)
    #contx_loss = tf.reduce_mean(l1-l2)
    #loss,resout = sess.run([contx_loss, pretrained_resnet.output],
    #        {in_expr[0]: latents, in_expr[1] : np.zeros((1,0))})
    #print(loss)
    #grads,_ = tf.gradients(0.01*pixelloss+0.1*contx_loss, in_expr)
    out_gpu_im = tflib.convert_images_to_uint8(out_gpu)
    target_expr_im = target_expr
    #target_expr = tflib.convert_images_from_uint8(tf.expand_dims(target_expr,0), nhwc_to_nchw=True)
    target_expr = tflib.convert_images_from_uint8(tf.expand_dims(target_expr,0), nhwc_to_nchw=True)
    #target_expr = tflib.convert_images_from_uint8(out_gpu_im)

    print("\n\n\n", out_gpu, target_expr, "\n\n\n")
    pixelloss = tf.nn.l2_loss(out_gpu - target_expr)
    grads,_ = tf.gradients(pixelloss, in_expr)
    #grads,_ = tf.gradients(contx_loss, in_expr)
    #exit()
    print("output tensors", out_gpu, target_expr)
    sess  = tf.get_default_session()
    for i in range(1000):
        print("gradient step",i)
        dx, ploss,out,t,og,te = sess.run([grads,pixelloss,out_gpu_im,target_expr_im,out_gpu, target_expr],
            {in_expr[0]: latents, in_expr[1] : np.zeros((1,0))})

        print("losses", ploss,ploss)
        #print(dx[0].shape, np.max(dx[0]))
        #latents = latents - ((0.00001/(1+0.1*i))*+dx[0])
        latents = latents - 0.00001* dx
        #np.save("og.np", og[0])
        #np.save("te.np", te[0])

    #writer.close()
        if i % 100 == 0:
           images = Gs.run(latents, None, truncation_psi=0.7, randomize_noise=False, output_transform=fmt)

           print("output tensors",len(out), out[0].shape, t.shape)
           PIL.Image.fromarray(images[0], 'RGB').save("results/{}_out.png".format(i))
           PIL.Image.fromarray(np.uint8(t), 'RGB').save("results/ground_target.png".format(i))
           PIL.Image.fromarray(out[0][0].transpose((1,2,0)), 'RGB').save("results/ground_out.png".format(i))

           print(out[0][0].shape)
           print(out[0][0].transpose((2,1,0)).shape)
           print(out[0][0].transpose((1,2,0)).shape)

           pass

if __name__ == "__main__":
    main()