from metrics import metric_base
import numpy as np
import tensorflow as tf
from tqdm import tqdm


class IL(metric_base.MetricBase):
    def __init__(self, num_images, interpolation_steps, step_size, **kwargs):
        super().__init__(**kwargs)
        self.num_images = num_images
        self.interpolation_steps = interpolation_steps
        self.step_size = step_size

    def _evaluate(self, Gs, Gs_kwargs, num_gpus):

        result = 0
        with  tf.device('/gpu:%d' % 0):
            for _ in tqdm(range(self.num_images)):
                latent = np.random.normal(0, 1, [1, Gs.input_shape[1]]).astype(np.float32)
                label_left = self._get_random_labels_np(1)
                rotation_offset = 108
                rotations = label_left[:, rotation_offset:rotation_offset + 8]
                rotation_index = np.argmax(rotations, axis=1)
                new_rotation_index = ((rotation_index + np.random.choice([-1, 1])) % 8)
                new_rotation = np.zeros([8], dtype=np.uint32)
                new_rotation[new_rotation_index] = 1
                new_rotation = new_rotation * np.expand_dims(
                    np.max(label_left[:, rotation_offset:rotation_offset + 8], axis=1).astype(np.uint32), axis=1)
                label_right = label_left.copy()
                label_right[:, rotation_offset:rotation_offset + 8] = new_rotation

                if 'mapping' in Gs.components:
                    dlatent_left = Gs.components.mapping.get_output_for(latent, label_left)
                    dlatent_right = Gs.components.mapping.get_output_for(latent, label_right)
                else:
                    dlatent = Gs.components.mapping_latent.get_output_for(latent)
                    dlabel_1 = Gs.components.mapping_label.get_output_for(label_left)
                    dlabel_2 = Gs.components.mapping_label.get_output_for(label_right)
                    if Gs.components.synthesis.input_shape[2] > 512:
                        dlatent_left = tf.concat([dlatent, dlabel_1], axis=-1)
                        dlatent_right = tf.concat([dlatent, dlabel_2], axis=-1)
                    else:
                        dlatent_left = dlatent + dlabel_1
                        dlatent_right = dlatent + dlabel_2

                # magnitudes = tf.sort(
                #     tf.clip_by_value(tf.random.normal(shape=[self.interpolation_steps - 2], mean=0.5, stddev=0.2), 0.0, 1.0))
                # zero = tf.zeros(shape=[1])
                # one = tf.ones(shape=[1])
                # magnitudes = tf.concat([zero, magnitudes, one], axis=0)
                magnitudes = tf.cast(tf.range(self.interpolation_steps) / self.interpolation_steps, tf.float32)

                x_axis = []
                y_axis = []
                for i in range(self.interpolation_steps):
                    x_axis.append(magnitudes[i])
                    prev_interpolate = dlatent_left * (1 - magnitudes[i] + self.step_size) + dlatent_right * (
                                magnitudes[i] - self.step_size)
                    dlatent_interpolate = prev_interpolate - dlatent_left * self.step_size + dlatent_right * self.step_size
                    pl_grads = tf.gradients(
                        Gs.components.synthesis.get_output_for(dlatent_interpolate, randomize_noise=False)
                        , [prev_interpolate])[0] - tf.gradients(
                        Gs.components.synthesis.get_output_for(prev_interpolate, randomize_noise=False)
                        , [prev_interpolate])[0]
                    y_axis.append(tf.sqrt(tf.reduce_mean(tf.reduce_sum(tf.square(pl_grads), axis=2), axis=1))[0])

                y_axis = tf.stack(y_axis)
                linear_diff = []
                y_axis_norm = (y_axis - tf.reduce_min(tf.stack([y_axis[0], y_axis[-1]]))) / tf.abs(
                    y_axis[0] - y_axis[-1])
                for i in range(self.interpolation_steps):
                    x = magnitudes[i]
                    f = y_axis_norm[0] + x * (y_axis_norm[-1] - y_axis_norm[0]) / x_axis[-1]
                    linear_diff.append(tf.square((y_axis_norm[i] - f) * (tf.abs(y_axis[0] - y_axis[-1])) / 1000))
                mean_linear_diff = tf.reduce_mean(tf.stack(linear_diff)).eval()
                result += mean_linear_diff
                print(mean_linear_diff)

            # self._report_progress(i, self.num_images)
            result = result / self.num_images
            self._report_result(result)
            print(result)

from metrics import metric_base
import numpy as np
import tensorflow as tf
from tqdm import tqdm


class IL(metric_base.MetricBase):
    def __init__(self, num_images, interpolation_steps, step_size, **kwargs):
        super().__init__(**kwargs)
        self.num_images = num_images
        self.interpolation_steps = interpolation_steps
        self.step_size = step_size

    def _evaluate(self, Gs, Gs_kwargs, num_gpus):

        result = 0
        with  tf.device('/gpu:%d' % 0):
            for _ in tqdm(range(self.num_images)):
                latent = np.random.normal(0, 1, [1, Gs.input_shape[1]]).astype(np.float32)
                label_left = self._get_random_labels_np(1)
                rotation_offset = 108
                rotations = label_left[:, rotation_offset:rotation_offset + 8]
                rotation_index = np.argmax(rotations, axis=1)
                new_rotation_index = ((rotation_index + np.random.choice([-1, 1])) % 8)
                new_rotation = np.zeros([8], dtype=np.uint32)
                new_rotation[new_rotation_index] = 1
                new_rotation = new_rotation * np.expand_dims(
                    np.max(label_left[:, rotation_offset:rotation_offset + 8], axis=1).astype(np.uint32), axis=1)
                label_right = label_left.copy()
                label_right[:, rotation_offset:rotation_offset + 8] = new_rotation

                if 'mapping' in Gs.components:
                    dlatent_left = Gs.components.mapping.get_output_for(latent, label_left)
                    dlatent_right = Gs.components.mapping.get_output_for(latent, label_right)
                else:
                    dlatent = Gs.components.mapping_latent.get_output_for(latent)
                    dlabel_1 = Gs.components.mapping_label.get_output_for(label_left)
                    dlabel_2 = Gs.components.mapping_label.get_output_for(label_right)
                    if Gs.components.synthesis.input_shape[2] > 512:
                        dlatent_left = tf.concat([dlatent, dlabel_1], axis=-1)
                        dlatent_right = tf.concat([dlatent, dlabel_2], axis=-1)
                    else:
                        dlatent_left = dlatent + dlabel_1
                        dlatent_right = dlatent + dlabel_2

                # magnitudes = tf.sort(
                #     tf.clip_by_value(tf.random.normal(shape=[self.interpolation_steps - 2], mean=0.5, stddev=0.2), 0.0, 1.0))
                # zero = tf.zeros(shape=[1])
                # one = tf.ones(shape=[1])
                # magnitudes = tf.concat([zero, magnitudes, one], axis=0)
                magnitudes = tf.cast(tf.range(self.interpolation_steps) / self.interpolation_steps, tf.float32)

                x_axis = []
                y_axis = []
                for i in range(self.interpolation_steps):
                    x_axis.append(magnitudes[i])
                    prev_interpolate = dlatent_left * (1 - magnitudes[i] + self.step_size) + dlatent_right * (
                                magnitudes[i] - self.step_size)
                    dlatent_interpolate = prev_interpolate - dlatent_left * self.step_size + dlatent_right * self.step_size
                    pl_grads = tf.gradients(
                        Gs.components.synthesis.get_output_for(dlatent_interpolate, randomize_noise=False)
                        , [prev_interpolate])[0] - tf.gradients(
                        Gs.components.synthesis.get_output_for(prev_interpolate, randomize_noise=False)
                        , [prev_interpolate])[0]
                    y_axis.append(tf.sqrt(tf.reduce_mean(tf.reduce_sum(tf.square(pl_grads), axis=2), axis=1))[0])

                y_axis = tf.stack(y_axis)
                linear_diff = []
                y_axis_norm = (y_axis - tf.reduce_min(tf.stack([y_axis[0], y_axis[-1]]))) / tf.abs(
                    y_axis[0] - y_axis[-1])
                for i in range(self.interpolation_steps):
                    x = magnitudes[i]
                    f = y_axis_norm[0] + x * (y_axis_norm[-1] - y_axis_norm[0]) / x_axis[-1]
                    linear_diff.append(tf.square((y_axis_norm[i] - f) * (tf.abs(y_axis[0] - y_axis[-1])) / 1000))
                mean_linear_diff = tf.reduce_mean(tf.stack(linear_diff)).eval()
                result += mean_linear_diff
                print(mean_linear_diff)

            # self._report_progress(i, self.num_images)
            result = result / self.num_images
            self._report_result(result)
            print(result)

# magnitudes = tf.sort( #     tf.clip_by_value(tf.random.normal(shape=[self.interpolation_steps - 2], mean=0.5, stddev=0.2), 0.0, 1.0)) # zero = tf.zeros(shape=[1]) # one = tf.ones(shape=[1]) # magnitudes = tf.concat([zero, magnitudes, one], axis=0) # self._report_progress(i, self.num_images)