import tensorflow as tf

from stylegan2.layers.embedding import PoseShapeEmbedding, GarmentEmbedding
from stylegan2.layers.dense import Dense
from stylegan2.layers.bias_act import BiasAct


class Mapping(tf.keras.layers.Layer):
    def __init__(self, w_dim, labels_dim, n_mapping, name, **kwargs):
        super(Mapping, self).__init__(name=name, **kwargs)
        self.w_dim = w_dim
        self.labels_dim = labels_dim
        self.n_mapping = n_mapping
        self.gain = 1.0
        self.lrmul = 0.01

        self.pose_shape_embedding = PoseShapeEmbedding()
        self.garment_embedding = GarmentEmbedding((self.w_dim * 2) - (216 + 11))
        # self.garment_embedding = GarmentEmbedding(self.w_dim)

        self.normalize = tf.keras.layers.Lambda(lambda x: x * tf.math.rsqrt(tf.reduce_mean(tf.square(x), axis=1, keepdims=True) + 1e-8))

        self.dense_layers = list()
        self.bias_act_layers = list()
        for ii in range(self.n_mapping):
            self.dense_layers.append(Dense(w_dim, gain=self.gain, lrmul=self.lrmul, name='dense_{:d}'.format(ii)))
            self.bias_act_layers.append(BiasAct(lrmul=self.lrmul, act='lrelu', name='bias_{:d}'.format(ii)))

    def call(self, inputs, training=None, mask=None):
        latents, conditions = inputs
        x = latents

        # embed label if any
        pose_shape_latent = self.pose_shape_embedding([conditions[0], conditions[1]])
        garment_latent = self.garment_embedding(conditions[2])
        x = tf.concat([x, pose_shape_latent, garment_latent], axis=1)
        # x = tf.concat([x, garment_latent], axis=1)

        # normalize inputs
        x = self.normalize(x)

        # apply mapping blocks
        for dense, apply_bias_act in zip(self.dense_layers, self.bias_act_layers):
            x = dense(x)
            x = apply_bias_act(x)

        return x

    def get_config(self):
        config = super(Mapping, self).get_config()
        config.update({
            'w_dim': self.w_dim,
            'labels_dim': self.labels_dim,
            'n_mapping': self.n_mapping,
            'gain': self.gain,
            'lrmul': self.lrmul,
        })
        return config