import numpy as np
import matplotlib.pyplot as plt
import dnnlib.tflib as tflib
import training.misc as misc
import os
from training import dataset
from tqdm import tqdm
import tensorflow as tf

tflib.init_tf()

fig, axes = plt.subplots(2, 4, subplot_kw=dict(projection='polar'), figsize=(16, 9))
fig.tight_layout(pad=3.0)

rotation_offset = 108
tfrecord_dir = '../../datasets/car_labels_new_rotation_oversample'
dset = dataset.TFRecordDataset(tfrecord_dir)
batch_size = 8
dset.configure(batch_size)

# network_dir = '../../results/00190-stylegan2-car_labels_new_rotation_oversample-2gpu-config-f-euclidean-interpolate-50-percent_continue_189-v7-256'
# network_dir = '../../results/00189-stylegan2-car_labels_new_rotation_oversample-2gpu-config-f-euclidean-v7-256-resume-00274'

network_dir = '../../results/00207-stylegan2-car_labels_v7_oversample_filter-2gpu-config-f-squared_euclidean_10_interpolate_50_percent_int_reg_remove_half_fl_fr_square-256'

file = sorted(os.listdir(network_dir))[-1]
G, D, Gs = misc.load_pkl(network_dir + '/' + file)
mean_dist_disc = []
mean_dist_gen = []

titles = [
    'Front Center',
    'Front Left',
    'Profile Left',
    'Rear Left',
    'Rear Center',
    'Rear Right',
    'Profile Right',
    'Front Right'
]

all_rotations = np.array([
    [1.0, 0.0],
    [0.7071, 0.7071],
    [0.0, 1.0],
    [-0.7071, 0.7071],
    [-1.0, 0.0],
    [-0.7071, -0.7071],
    [0.0, -1.0],
    [0.7071, -0.7071],
    [0., 0.]
],
    dtype=np.float32
)
all_rotations = np.tile(np.expand_dims(all_rotations, axis=0), [batch_size, 1, 1])
all_label_rotations = [[], [], [], [], [], [], [], [], []]
all_fake_rotations = [[], [], [], [], [], [], [], [], []]
all_pred_rotations = [[], [], [], [], [], [], [], [], []]

label_placeholder = tf.placeholder(tf.float32, shape=G.input_shapes[1])
image_placeholder = tf.placeholder(tf.float32, shape=D.input_shapes[0])

latents = tf.random_normal(shape=[batch_size] + G.input_shapes[0][1:])
fake_images = G.get_output_for(latents, label_placeholder, is_training=True)
real_scores = D.get_output_for(image_placeholder, label_placeholder, is_training=True)
fake_scores = D.get_output_for(fake_images, label_placeholder, is_training=True)

for i in tqdm(range(30)):
    images, labels = dset.get_minibatch_np(batch_size)
    images = misc.adjust_dynamic_range(images, [0, 255], [-1, 1])
    fake_images_out, real_scores_out, fake_scores_out = tflib.run([fake_images, real_scores, fake_scores],
                                                                  feed_dict={
                                                                      label_placeholder: labels,
                                                                      image_placeholder: images
                                                                  })

    labels_rotation = labels[:, rotation_offset:rotation_offset + 2]
    fake_rotations = fake_scores_out[:, rotation_offset:rotation_offset + 2]
    pred_rotations = real_scores_out[:, rotation_offset:rotation_offset + 2]

    labels_rotation_tile = np.tile(np.expand_dims(labels_rotation, axis=1), [1, 9, 1])
    distances = np.linalg.norm(all_rotations - labels_rotation_tile, axis=-1)
    rotation_index = np.argmin(distances, axis=-1)
    print(rotation_index)
    for j in range(batch_size):
        all_label_rotations[rotation_index[j]].append(labels_rotation[j])
        all_fake_rotations[rotation_index[j]].append(fake_rotations[j])
        all_pred_rotations[rotation_index[j]].append(pred_rotations[j])

index = 0

for i in range(2):
    for j in range(4):
        mean_disc = 0
        mean_gen = 0
        for k in range(len(all_label_rotations[index])):
            ax = axes[i, j]
            ax.set_rscale('linear')
            current_label = np.array(all_label_rotations[index][k])
            current_pred = np.array(all_pred_rotations[index][k])
            current_fake = np.array(all_fake_rotations[index][k])
            mean_disc += np.linalg.norm(current_label - current_pred, ord=2)
            mean_gen += np.linalg.norm(current_label - current_fake, ord=2)

            abs_label_real = [np.linalg.norm(current_label, ord=2), np.linalg.norm(current_pred, ord=2)]
            angle_label_real = [np.arctan2(current_label[1], current_label[0]),
                                np.arctan2(current_pred[1], current_pred[0])]
            ax.plot(angle_label_real, abs_label_real, c='green')
            ax.scatter(angle_label_real, abs_label_real, c='green')

            abs_label_fake = [np.linalg.norm(current_label, ord=2), np.linalg.norm(current_fake, ord=2)]
            angle_label_fake = [np.arctan2(current_label[1], current_label[0]),
                                np.arctan2(current_fake[1], current_fake[0])]
            ax.plot(angle_label_fake, abs_label_fake, c='red')
            ax.scatter(angle_label_fake, abs_label_fake, c='red')

            # abs_zero_fake = [0, np.linalg.norm(current_fake, ord=2)]
            # angle_zero_fake = [0, np.arctan2(current_fake[1], current_fake[0])]
            # ax.plot(angle_zero_fake, abs_zero_fake, c='orange')
            # ax.scatter(angle_zero_fake, abs_zero_fake, c='orange')
        mean_disc /= len(all_label_rotations[index])
        mean_gen /= len(all_label_rotations[index])
        mean_dist_disc.append(mean_disc)
        mean_dist_gen.append(mean_gen)
        index += 1

index = 0
for i in range(2):
    for j in range(4):
        ax = axes[i, j]
        ax.title.set_text(titles[index] + '\nmean_dist_disc: ' + str(np.round(mean_dist_disc[index], 4)) + '\nmean_dist_gen: ' + str(np.round(mean_dist_gen[index], 4)))
        index += 1

plt.savefig("plot_circle_multiple_207.png")
plt.show()
