import tensorflow as tf

from convnext.models.convnext_tf import ConvNeXt

def rodrigues(r):
    theta = tf.norm(r, axis=(1, 2), keepdims=True)
    theta = tf.maximum(theta, 1e-7)
    r_hat = r / theta
    z_stick = tf.zeros(theta.get_shape().as_list()[0], dtype=tf.float32)
    m = tf.stack(
        (z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick,
        -r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick), axis=1)
    m = tf.reshape(m, (-1, 3, 3))
    i_cube = tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros(
        (theta.get_shape().as_list()[0], 3, 3), dtype=tf.float32)
    # R = cos * i_cube + (1 - cos) * dot + tf.sin(theta) * m
    R = i_cube + m * tf.sin(theta) + tf.matmul(m, m) * (1 - tf.cos(theta))
    return R


def run_rodrigues(pose):
    batch_size = pose.shape[0]
    R_out = rodrigues(tf.reshape(pose, (-1, 1, 3)))
    R_out = tf.reshape(R_out, (batch_size, -1, 3, 3))
    return R_out


# if __name__ == '__main__':
#     x = tf.ones((4, 72))
#     y = run_rodrigues(x)
#     print()
#     pass


class ParameterEmbedding(tf.keras.layers.Layer):
    def __init__(self, embed_dim, **kwargs):
        super(ParameterEmbedding, self).__init__(**kwargs)
        self.embed_dim = embed_dim

        self.dense_1 = tf.keras.layers.Dense(128, activation='tanh', name='dense_1')
        self.bn_1 = tf.keras.layers.BatchNormalization()
        self.dense_2 = tf.keras.layers.Dense(512, activation='tanh', name='dense_2')
        self.bn_2 = tf.keras.layers.BatchNormalization()
        self.dense_3 = tf.keras.layers.Dense(self.embed_dim, activation='tanh', name='dense_3')

    def call(self, inputs, training=None, mask=None):
        x = inputs
        x = self.dense_1(x)
        x = self.bn_1(x)
        x = self.dense_2(x)
        x = self.bn_2(x)
        x = self.dense_3(x)
        return x

    def get_config(self):
        config = super(ParameterEmbedding, self).get_config()
        config.update({
            'embed_dim': self.embed_dim,
        })
        return config


class PoseShapeEmbedding(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        super(PoseShapeEmbedding, self).__init__(**kwargs)

        # self.pose_embedding = ParameterEmbedding(int(self.embed_dim/2))
        # self.shape_embedding = ParameterEmbedding(int(self.embed_dim/2))
        self.flatten = tf.keras.layers.Flatten()

    def call(self, inputs, training=None, mask=None):
        pose, shape = inputs
        rotations = run_rodrigues(pose)
        rotations = self.flatten(rotations)
        # pose_latent = self.pose_embedding(rotations, training)
        # shape_latent = self.shape_embedding(shape, training)
        return tf.concat((rotations, shape), axis=-1)

    def get_config(self):
        config = super(PoseShapeEmbedding, self).get_config()
        config.update({
            'embed_dim': 216 + 11,
        })
        return config


class GarmentEmbedding(tf.keras.layers.Layer):
    def __init__(self, embed_dim, **kwargs):
        super(GarmentEmbedding, self).__init__(**kwargs)
        self.embed_dim = embed_dim

        # TODO: channels_first?

        self.encoder = ConvNeXt(depths=[3, 3, 27, 3], dims=[96, 192, 256, self.embed_dim], include_top=False)
        # self.encoder(tf.keras.Input(shape=(*input_shape, 3)))
        self.pooling = tf.keras.layers.GlobalAveragePooling2D()

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

        x = tf.transpose(x, [0,2,3,1])
        x = self.encoder(x, training=training)
        x = self.pooling(x)

        return x

    def get_config(self):
        config = super(GarmentEmbedding, self).get_config()
        config.update({
            'embed_dim': self.embed_dim,
        })
        return config


# class LabelEmbedding(tf.keras.layers.Layer):
#     def __init__(self, embed_dim, **kwargs):
#         super(LabelEmbedding, self).__init__(**kwargs)
#         self.embed_dim = embed_dim

#     def build(self, input_shape):
#         weight_shape = [input_shape[1], self.embed_dim]
#         # tf 1.15 mean(0.0), std(1.0) default value of tf.initializers.random_normal()
#         w_init = tf.random.normal(shape=weight_shape, mean=0.0, stddev=1.0)
#         self.w = tf.Variable(w_init, name='w', trainable=True)

#     def call(self, inputs, training=None, mask=None):
#         x = tf.matmul(inputs, self.w)
#         return x

#     def get_config(self):
#         config = super(LabelEmbedding, self).get_config()
#         config.update({
#             'embed_dim': self.embed_dim,
#         })
#         return config