from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
import sys
import os
from datetime import datetime
from tqdm import tqdm
from PIL import Image
import imageio

import warnings
warnings.filterwarnings('ignore', category=FutureWarning)

# ffmpeg installation location, for creating videos
plt.rcParams['animation.ffmpeg_path'] = str(Path.home() / "Documents/dev_tools/ffmpeg-20190623-ffa64a4-win64-static/bin/ffmpeg.exe")

%load_ext autoreload
%autoreload 2

# StyleGAN Utils
from stylegan_utils import load_network, gen_image_fun, synth_image_fun, create_video, map_latents

import dnnlib
import dataset_tool
import run_projector
import projector
import training.dataset
import training.misc

# Specific of encoder repos, comment out if not needed
#from encoder.perceptual_model import PerceptualModel
#from encoder.generator_model import Generator

# Data Science Utils
sys.path.append(os.path.join(os.pardir, 'data-science-learning'))

from ds_utils import generative_utils

res_dir = Path.home() / 'Documents/generated_data/stylegan'

"""
Explanation: <h1>Table of Contents<span class="tocSkip"></span></h1>
<div class="toc"><ul class="toc-item"><li><span><a href="#Load-Network" data-toc-modified-id="Load-Network-1"><span class="toc-item-num">1&nbsp;&nbsp;</span>Load Network</a></span></li><li><span><a href="#Project" data-toc-modified-id="Project-2"><span class="toc-item-num">2&nbsp;&nbsp;</span>Project</a></span></li><li><span><a href="#Encode" data-toc-modified-id="Encode-3"><span class="toc-item-num">3&nbsp;&nbsp;</span>Encode</a></span></li><li><span><a href="#Generate-Images" data-toc-modified-id="Generate-Images-4"><span class="toc-item-num">4&nbsp;&nbsp;</span>Generate Images</a></span><ul class="toc-item"><li><span><a href="#generate-some-random-samples-and-save-to-disk" data-toc-modified-id="generate-some-random-samples-and-save-to-disk-4.1"><span class="toc-item-num">4.1&nbsp;&nbsp;</span>generate some random samples and save to disk</a></span></li></ul></li><li><span><a href="#Projected-Latent-Initialization" data-toc-modified-id="Projected-Latent-Initialization-5"><span class="toc-item-num">5&nbsp;&nbsp;</span>Projected Latent Initialization</a></span><ul class="toc-item"><li><span><a href="#Load-latents-from-other-network-and-generate-with-current" data-toc-modified-id="Load-latents-from-other-network-and-generate-with-current-5.1"><span class="toc-item-num">5.1&nbsp;&nbsp;</span>Load latents from other network and generate with current</a></span></li></ul></li></ul></div>

Project/embed real images into StyleGANv2 latent space.
End of explanation
"""


MODELS_DIR = Path("C:/Users/User/Documents/models/stylegan2")
MODEL_NAME = 'drawing2_1024'
SNAPSHOT_NAME = 'network-snapshot-002048'

Gs, Gs_kwargs, noise_vars = load_network(str(MODELS_DIR / MODEL_NAME / SNAPSHOT_NAME) + '.pkl')

Z_SIZE = Gs.input_shape[1:][0]
IMG_SIZE = Gs.output_shape[2:]
IMG_SIZE

"""
Explanation: Load Network
End of explanation
"""


def project_images(images_dir, tfrecord_dir, data_dir, num_steps, num_snapshots, pure_projector=False,
                  lr=0.1, dlatent_init_dir=None):
    # setup projector
    print('Setting up projector')
    proj = projector.Projector(num_steps=num_steps, pure_projector=pure_projector, initial_learning_rate=lr)
    proj.set_network(Gs)
    
    # generate tfrecords
    images_paths = dataset_tool.create_from_images(str(tfrecord_dir), str(images_dir), False)
    nb_images = len(images_paths)

    # loading images from tfrecords
    dataset_obj = training.dataset.load_dataset(data_dir=str(data_dir), tfrecord_dir=tfrecord_dir, 
                                                max_label_size=0, verbose=True, repeat=False, shuffle_mb=0)
    assert dataset_obj.shape == Gs.output_shape[1:]
    
    # project all loaded images
    print('=======================')
    for image_idx in tqdm(range(nb_images)):
        print(f'Projecting image {image_idx}/{nb_images}')
        
        images, _labels = dataset_obj.get_minibatch_np(1)
        images = training.misc.adjust_dynamic_range(images, [0, 255], [-1, 1])
        
        # loading init dlatent, if given
        if dlatent_init_dir is not None:
            dlatent_init = np.load(dlatent_init_dir / (Path(images_paths[image_idx]).stem + '.npy'))
        else:
            dlatent_init = None
        
        run_path = data_dir / f'out_{image_idx}'
        run_path.mkdir()
        run_projector.project_image(proj, targets=images, 
                                    png_prefix=dnnlib.make_run_dir_path(str(run_path / 'image_')), 
                                    num_snapshots=num_snapshots,
                                    dlatent_init=dlatent_init)

data_dir = res_dir / 'projection' / MODEL_NAME / SNAPSHOT_NAME / datetime.now().strftime("%Y%m%d_%H%M%S") # where the projections results will be saved
images_dir = Path.home() / 'Documents/generated_data/'

tfrecord_dir = data_dir / 'tfrecords'
project_images(images_dir=images_dir, tfrecord_dir=tfrecord_dir, data_dir=data_dir, 
               num_steps=1000, num_snapshots=10, pure_projector=False,
               dlatent_init_dir=None)#data_dir.parent / 'latent_init')

create_video(data_dir / 'out', res_dir / 'projection' / 'out_{}.mp4'.format(''))

"""
Explanation: Project
End of explanation
"""


from encoder.perceptual_model import PerceptualModel
from encoder.generator_model import Generator

import pretrained_networks
_G, _D, Gs = pretrained_networks.load_networks((str(MODELS_DIR / MODEL_NAME / SNAPSHOT_NAME) + '.pkl'))

BATCH_SIZE = 1
PERCEPTUAL_MODEL_IMG_SIZE = 256

# setup utils generator and perceptual model
generator = Generator(Gs, BATCH_SIZE, randomize_noise=False)
perceptual_model = PerceptualModel(PERCEPTUAL_MODEL_IMG_SIZE, layer=9, batch_size=BATCH_SIZE)
perceptual_model.build_perceptual_model(generator, _D)

def split_to_batches(l, n):
    for i in range(0, len(l), n):
        yield l[i:i + n]

def encode_images(images_dir, data_dir, iterations, num_snapshots, learning_rate=1.):
    # collect images
    images_paths = [str(img) for img in images_dir.glob('*')][:3]
    
    run_path = data_dir
    run_path.mkdir()
    
    snapshot_steps = set(iterations - np.linspace(0, iterations, num_snapshots, endpoint=False, dtype=int))
    
    # project all loaded images
    for images_batch in tqdm(split_to_batches(images_paths, BATCH_SIZE), total=len(images_paths)//BATCH_SIZE):
        cur_step = 0
        images_names = [os.path.splitext(os.path.basename(img_path))[0] for img_path in images_batch]

        perceptual_model.set_reference_images(images_batch)
        optimizer = perceptual_model.optimize(generator.dlatent_variable, 
                                       iterations=iterations)
        pbar = tqdm(optimizer, leave=False, mininterval=9, total=iterations)
        for loss in pbar:
            cur_step += 1
            if cur_step in snapshot_steps:
                #print(' '.join(images_names), ' loss:', loss)

                # generate images from found dlatents and save them
                generated_images = generator.generate_images()
                generated_dlatents = generator.get_dlatents()
                for img_array, dlatent, img_name in zip(generated_images, generated_dlatents, images_names):
                    img = Image.fromarray(img_array, 'RGB')
                    img.save(str(run_path / f'{cur_step}_{img_name}.png'), 'PNG')
                    np.save(str(run_path / f'{cur_step}_{img_name}.npy'), dlatent)

        generator.reset_dlatents()

data_dir = res_dir / 'projection' / MODEL_NAME / SNAPSHOT_NAME / datetime.now().strftime("%Y%m%d_%H%M%S") # where the projections results will be saved
images_dir = Path.home() / 'Documents/generated_data/'

encode_images(images_dir=images_dir, data_dir=data_dir, iterations=1000,  num_snapshots=10)

"""
Explanation: Encode
This does not use the official StyleGAN v2 projector, but instead relies on the direct encoder setup used by the community for v1.
End of explanation
"""


target_latents = np.random.randn(18, Z_SIZE)

img = gen_image_fun(Gs, target_latents, Gs_kwargs, noise_vars, truncation_psi=0.5)
plt.imshow(img)

img = synth_image_fun(Gs, target_latents[np.newaxis,:,:], Gs_kwargs, randomize_noise=True)
plt.imshow(img)

"""
Explanation: Generate Images
End of explanation
"""


zlatents = np.random.randn(30, Z_SIZE)
truncation_psi = 0.5

data_dir = res_dir / 'projection' / MODEL_NAME / SNAPSHOT_NAME / 'rand_gen' / datetime.now().strftime("%Y%m%d_%H%M%S")
data_dir.mkdir(parents=True, exist_ok=False)

dlatents = map_latents(Gs, zlatents, truncation_psi)
for i, dl in enumerate(dlatents):
    img  = synth_image_fun(Gs, dl[np.newaxis,:,:], Gs_kwargs, randomize_noise=True)
    img = Image.fromarray(img, 'RGB')
    img.save(str(data_dir / f'{i}.png'), 'PNG')
    np.save(str(data_dir / f'{i}.npy'), dl)

"""
Explanation: generate some random samples and save to disk
End of explanation
"""


from keras.models import load_model

resnet = load_model(MODELS_DIR / MODEL_NAME / 'resnet' / 'finetuned_resnet.h5')
resnet.summary()

resnet_img_size = (512, 512)

resnet.summary()

target_img = Image.open("")
target_img = target_img.resize(resnet_img_size)
plt.imshow(target_img)

predicted_latent = resnet.predict(np.array(target_img)[np.newaxis,:])

img = synth_image_fun(Gs, predicted_latent, Gs_kwargs, randomize_noise=True)
plt.imshow(img)

def generate_latent_init(resnet_path, images_dir, out_dir):
    # setup projector
    print('Setting up resnet model')
    resnet = load_model(resnet_path)
    resnet_img_size = (512, 512)
    
    # project all images
    print('=======================')
    all_images = list(images_dir.glob('*.png')) + list(images_dir.glob('*.jpg'))
    for img_path in tqdm(all_images):
        target_img = Image.open(str(img_path))
        target_img = target_img.resize(resnet_img_size)
        
        predicted_latent = resnet.predict(np.array(target_img)[np.newaxis,:])[0]
        
        out_dir.mkdir(parents=True, exist_ok=True)
        np.save(out_dir / (img_path.stem + '.npy'), predicted_latent)

latent_init_dir = res_dir / 'projection' / MODEL_NAME / SNAPSHOT_NAME / 'latent_init'
images_dir = Path.home() / 'Documents/generated_data/'

generate_latent_init(resnet_path=MODELS_DIR / MODEL_NAME / 'resnet' / 'finetuned_resnet.h5',
               images_dir=images_dir, out_dir=latent_init_dir)

"""
Explanation: Projected Latent Initialization
Test network used to learn an initial mapping from an image to the intermediate StyleGAN latent
End of explanation
"""


load_dir = res_dir / 'projection' / 'original_ffhq' / 'stylegan2-ffhq-config-f'
out_dir = res_dir / 'projection' / MODEL_NAME / SNAPSHOT_NAME / 'dlatent_mix' / datetime.now().strftime("%Y%m%d_%H%M%S")

out_dir.mkdir(parents=True, exist_ok=False)

for i, dl_path in enumerate(load_dir.glob('*.npy')):
    dl = np.load(dl_path)

    img = synth_image_fun(Gs, dl[np.newaxis,:,:], Gs_kwargs, randomize_noise=False)
    other_img = imageio.imread(str(dl_path).replace('npy', 'png'))
    
    w = h = img.shape[0]
    canvas = Image.new('RGBA', (w*2,h), 'white')
    canvas.paste(Image.fromarray(img), (0, 0))
    canvas.paste(Image.fromarray(other_img), (w, 0))
    
    canvas.save(str(out_dir / f'{i}.png'), 'PNG')

"""
Explanation: Load latents from other network and generate with current
End of explanation
"""


