import tensorflow as tf

import os
import sys
import cv2
import numpy as np
from PIL import Image
from dataset import DatasetSimple
from load_models import load_generator
from tf_utils import allow_memory_growth

os.environ["CUDA_VISIBLE_DEVICES"] = "3"

def m2uv_array(m2uv):
    result = np.zeros((14475, 2), dtype=np.int32)
    for k, v in m2uv.items():
        if len(v) == 2:
            result[k] = list(v)
        elif len(v) == 1:
            x = v.pop()
            result[k] = [x, x]
        else:
            assert False

    return result


def uv_to_mesh(uv):
    uv_res = 256
    m2uv = np.load('MeshUtils/m2uv.npy', allow_pickle=True).item()
    m2uv = m2uv_array(m2uv)
    uv_map = np.load('Meshutils/stretched_uv_map_3.npy')
    uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32)

    mesh = np.zeros((14475, 3))
    uv_pixels = uv_map_discreet[m2uv]
    # counter = np.ones((14475))
    # counter[m2uv[:, 0] == m2uv[:, 1]] = 2
    mesh = np.mean(uv[uv_pixels[:, :, 0], uv_pixels[:, :, 1], :], axis=1) # / counter[:, np.newaxis]
    return mesh


def inference(ckpt_dir, use_custom_cuda, res, folder, data_folder, samples=50):
    # create generator
    resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024]
    featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32]
    filter_index = resolutions.index(res)
    g_params = {
        'z_dim': 256,
        'w_dim': 256,
        'labels_dim': 0,
        'n_mapping': 8,
        'resolutions': resolutions[:filter_index + 1],
        'featuremaps': featuremaps[:filter_index + 1],
    }
    generator = load_generator(g_params, is_g_clone=True, ckpt_dir=ckpt_dir, custom_cuda=use_custom_cuda)
    dataset = DatasetSimple(data_folder).cache().batch(1).repeat()

    os.makedirs(os.path.join(folder), exist_ok=True)

    i = 0
    for batch in dataset:
        print(i)
        # generate image
        fake_images = generator([tf.random.normal(shape=[1, g_params['z_dim']]),
                                batch[1]],
                                training=False, truncation_psi=0.5)

        np.savez_compressed(os.path.join(folder, str(i) + '.npz'), uv=tf.squeeze(fake_images), template=tf.squeeze(batch[1][2]), body=tf.squeeze(batch[3]), mask=tf.squeeze(batch[4]))
        img = np.transpose(np.round((fake_images + 1) * 128)[0], [1,2,0])
        cv2.imwrite(os.path.join(folder, str(i) + '.jpg'), img)
        i += 1

        if i > samples:
            break

    return


def main():
    allow_memory_growth()

    checkpoints = [
        {
            'res': 512,
            'ckpt_dir': 'models/stylegan2-test-8-512x512',
            'use_custom_cuda': True,
            'folder': 'output/test-8-196000',
            'data_folder': '/home/workspace/data/dataset/stylegan_new/train'
        },
    ]

    for run_item in checkpoints:
        res = run_item['res']
        ckpt_dir = run_item['ckpt_dir']
        use_custom_cuda = run_item['use_custom_cuda']
        folder = run_item['folder']
        message = f'{res}x{res} with custom cuda' if use_custom_cuda else f'{res}x{res} without custom cuda'
        data_folder = run_item['data_folder']

        print(message)
        inference(ckpt_dir, use_custom_cuda, res, folder, data_folder)
    return


if __name__ == '__main__':
    main()