#  Tencent is pleased to support the open source community by making GNES available.
#
#  Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved.
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#  http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

import math

import tensorflow as tf
import tensorflow.contrib.slim as slim


class NetFV():
    def __init__(self, feature_size,
                 cluster_size,
                 vocab_size,
                 method='netvlad',
                 use_length=False,
                 input_size=None,
                 use_2nd_label=False,
                 vocab_size_2=None,
                 add_batch_norm=True,
                 is_training=False,
                 use_weights=True,
                 save_dir=None,
                 multitask_method=None,
                 l2_penalty=1e-6):
        if input_size == None:
            self.input_size = feature_size
        else:
            self.input_size = input_size
        self.feature_size = feature_size
        self.use_length = use_length
        self.is_training = is_training
        self.vocab_size = vocab_size
        self.use_2nd_label = use_2nd_label
        self.vocab_size_2 = vocab_size_2
        self.add_batch_norm = add_batch_norm
        self.cluster_size = cluster_size
        self.use_weights = use_weights
        self.l2_penalty = l2_penalty
        self.method = method
        self.multitask_method = multitask_method
        self.build_model()
        self.build_loss()

    @staticmethod
    def rand_init(feature_size):
        return tf.random_normal_initializer(stddev=1 / math.sqrt(feature_size))

    def build_model(self):
        self.feeds = tf.placeholder(tf.float32, [None, None, self.input_size])
        self.feeds_length = tf.placeholder(tf.int32, [None])

        self.inputs = tf.layers.dense(self.feeds, self.feature_size)

        self.weights = tf.placeholder(tf.float32, [None, self.vocab_size])
        self.max_frames = tf.shape(self.inputs)[1]
        self.seq_length = tf.cast(tf.sequence_mask(self.feeds_length,
                                                   self.max_frames), tf.float32)
        if self.method == 'fvnet':
            self.build_fvnet()
        elif self.method == 'netvlad':
            self.build_netvlad()
        elif self.method == 'pooling':
            self.build_pooling()

    def build_pooling(self):
        self.repre = tf.layers.dense(self.inputs, self.feature_size)
        self.repre = tf.reduce_max(self.repre, axis=1)

    def build_fvnet(self):
        reshaped_input = tf.reshape(self.inputs, [-1, self.feature_size])
        cluster_weights = tf.get_variable("cluster_weights",
                                          [self.feature_size, self.cluster_size],
                                          initializer=NetFV.rand_init(self.feature_size))

        covar_weights = tf.get_variable("covar_weights",
                                        [self.feature_size, self.cluster_size],
                                        initializer=NetFV.rand_init(self.feature_size))

        covar_weights = tf.square(covar_weights)
        eps = tf.constant([1e-6])
        covar_weights = tf.add(covar_weights, eps)

        tf.summary.histogram("cluster_weights", cluster_weights)
        activation = tf.matmul(reshaped_input, cluster_weights)
        if self.add_batch_norm:
            activation = slim.batch_norm(activation,
                                         center=True,
                                         scale=True,
                                         is_training=self.is_training,
                                         scope="cluster_bn")
        else:
            cluster_biases = tf.get_variable("cluster_biases",
                                             [self.cluster_size],
                                             initializer=NetFV.rand_init(self.feature_size))
            tf.summary.histogram("cluster_biases", cluster_biases)
            activation += cluster_biases

        activation = tf.nn.softmax(activation)
        tf.summary.histogram("cluster_output", activation)

        activation = tf.reshape(activation, [-1, self.max_frames, self.cluster_size])

        a_sum = tf.reduce_sum(activation, -2, keepdims=True)

        cluster_weights2 = tf.scalar_mul(0.01, cluster_weights)

        a = tf.multiply(a_sum, cluster_weights2)

        activation = tf.transpose(activation, perm=[0, 2, 1])

        reshaped_input = tf.reshape(reshaped_input,
                                    [-1, self.max_frames, self.feature_size])
        fv1 = tf.matmul(activation, reshaped_input)

        fv1 = tf.transpose(fv1, perm=[0, 2, 1])

        # computing second order FV
        a2 = tf.multiply(a_sum, tf.square(cluster_weights2))

        b2 = tf.multiply(fv1, cluster_weights2)
        fv2 = tf.matmul(activation, tf.square(reshaped_input))

        fv2 = tf.transpose(fv2, perm=[0, 2, 1])
        fv2 = tf.add_n([a2, fv2, tf.scalar_mul(-2, b2)])

        fv2 = tf.divide(fv2, tf.square(covar_weights))
        fv2 = tf.subtract(fv2, a_sum)

        fv2 = tf.reshape(fv2, [-1, self.cluster_size * self.feature_size])
        fv2 = tf.nn.l2_normalize(fv2, 1)
        fv2 = tf.reshape(fv2, [-1, self.cluster_size * self.feature_size])
        fv2 = tf.nn.l2_normalize(fv2, 1)

        fv1 = tf.subtract(fv1, a)
        fv1 = tf.divide(fv1, covar_weights)
        fv1 = tf.nn.l2_normalize(fv1, 1)
        fv1 = tf.reshape(fv1, [-1, self.cluster_size * self.feature_size])
        fv1 = tf.nn.l2_normalize(fv1, 1)

        self.repre = tf.concat([fv1, fv2], 1)
        self.repre = tf.layers.dense(self.repre, self.feature_size)

    def build_netvlad(self):
        reshaped_input = tf.reshape(self.inputs, [-1, self.feature_size])
        cluster_weights = tf.get_variable("cluster_weights",
                                          [self.feature_size, self.cluster_size],
                                          initializer=NetFV.rand_init(self.feature_size))
        activation = tf.matmul(reshaped_input, cluster_weights)
        if self.add_batch_norm:
            activation = slim.batch_norm(activation,
                                         center=True,
                                         scale=True,
                                         is_training=self.is_training,
                                         scope="cluster_bn")
        else:
            cluster_biases = tf.get_variable("cluster_biases",
                                             [self.cluster_size],
                                             initializer=NetFV.rand_init(self.feature_size))
            activation += cluster_biases
        activation = tf.nn.softmax(activation)
        activation = tf.reshape(activation, [-1, self.max_frames, self.cluster_size])
        if self.use_length:
            activation *= tf.reshape(self.seq_length, [-1, self.max_frames, 1])

        a_sum = tf.reduce_sum(activation, -2, keep_dims=True)

        if self.use_length:
            a_sum = a_sum / tf.cast(tf.reshape(self.feeds_length, [-1, 1, 1]), tf.float32)

        cluster_weights2 = tf.get_variable("cluster_weights2",
                                           [1, self.feature_size, self.cluster_size],
                                           initializer=NetFV.rand_init(self.feature_size))

        a = tf.multiply(a_sum, cluster_weights2)
        activation = tf.transpose(activation, perm=[0, 2, 1])

        reshaped_input = tf.reshape(reshaped_input,
                                    [-1, self.max_frames, self.feature_size])
        vlad = tf.matmul(activation, reshaped_input)
        if self.use_length:
            vlad = vlad / tf.cast(tf.reshape(self.feeds_length, [-1, 1, 1]), tf.float32)
        vlad = tf.transpose(vlad, perm=[0, 2, 1])
        vlad = tf.subtract(vlad, a)

        vlad = tf.nn.l2_normalize(vlad, 1)

        vlad = tf.reshape(vlad, [-1, self.cluster_size * self.feature_size])
        vlad = tf.nn.l2_normalize(vlad, 1)
        self.repre = vlad

    def build_loss(self):
        self.probabilities = tf.layers.dense(self.repre,
                                             self.vocab_size,
                                             activation=tf.nn.tanh)
        self.probabilities = tf.layers.dense(self.probabilities, self.vocab_size)
        self.probabilities = tf.nn.softmax(self.probabilities)

        self.label = tf.placeholder(tf.int32, [None, self.vocab_size])
        logits = tf.cast(self.label, tf.float32)
        if self.use_weights:
            logits = logits * self.weights
        self.loss = - tf.log(tf.reduce_sum(logits * self.probabilities, axis=1) + 1e-9)
        self.loss = tf.reduce_mean(self.loss)
        self.pred = tf.argmax(self.probabilities, 1)
        self.avg_diff = tf.cast(tf.equal(tf.argmax(self.label, 1), self.pred), tf.float32)
        self.avg_diff = tf.reduce_mean(self.avg_diff)

        # add 2nd layer labels
        if self.use_2nd_label:
            self.label_2 = tf.placeholder(tf.int32, [None, self.vocab_size_2])
            logits2 = tf.cast(self.label_2, tf.float32)

            if self.multitask_method is None:
                self.probabilities2 = tf.layers.dense(self.repre,
                                                      self.vocab_size_2,
                                                      activation=tf.nn.tanh)
                self.probabilities2 = tf.layers.dense(self.probabilities2, self.vocab_size_2)
                self.probabilities2 = tf.nn.softmax(self.probabilities2)

            elif self.multitask_method == 'Attention':
                self.x = tf.get_variable('emb',
                                         shape=[self.vocab_size, self.feature_size],
                                         dtype=tf.float32,
                                         initializer=NetFV.rand_init(self.feature_size))
                self.emb_label = tf.matmul(self.probabilities, self.x)
                self.emb_concat = tf.concat([self.emb_label, self.repre], axis=1)
                self.probabilities2 = tf.layers.dense(self.emb_concat,
                                                      self.vocab_size_2,
                                                      activation=tf.nn.tanh)
                self.probabilities2 = tf.layers.dense(self.probabilities2,
                                                      self.vocab_size_2)
                self.probabilities2 = tf.nn.softmax(self.probabilities2)

            self.loss += tf.reduce_mean(-tf.log(
                tf.reduce_sum(logits2 * self.probabilities2, axis=1) + 1e-9))
            self.pred2 = tf.argmax(self.probabilities2, 1)
            self.avg_diff2 = tf.cast(tf.equal(tf.argmax(self.label_2, 1), self.pred2), tf.float32)
            self.avg_diff2 = tf.reduce_mean(self.avg_diff2)

        self.optimizer = tf.train.AdamOptimizer(learning_rate=0.0005,
                                                epsilon=1e-08,
                                                name='adam')
        self.train_op = slim.learning.create_train_op(self.loss, self.optimizer)
        self.eval_res = {'loss': self.loss, 'avg_diff': self.avg_diff}
        if self.use_2nd_label:
            self.eval_res['avg_diff2'] = self.avg_diff2
