"""Copyright 2020 Google LLC

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
    https://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.

Defines the architecture of the Video Classifier.
"""
import math
import tensorflow as tf

class NetVLAD(tf.keras.layers.Layer):
  """Applies NetVLAD to the input.
  
  Args:
    num_clusters: The number of clusters to use.
    input_shape: 3D tensor denoting the input shape of the NetVLAD layer.
  
  Input Shape:
    3D tensor with shape: `(batch_size, time, feature_dim)`.
  Output shape:
    2D tensor with shape: `(batch_size, feature_dim * num_clusters)`.
  """
  def __init__(self, num_clusters, input_shape, **kwargs):
    super().__init__(**kwargs)
    if num_clusters <= 0:
      raise ValueError("`num_clusters` must be greater than 1: %i" % num_clusters)
    self.num_clusters = num_clusters
    feature_dim = input_shape[-1]
    if not isinstance(feature_dim, int):
      feature_dim = feature_dim.value
    self.fc = tf.keras.layers.Dense(
      units=self.num_clusters,
      activation=tf.nn.softmax,
      kernel_regularizer=tf.keras.regularizers.l2(1e-5),
      name="vlad_fc" + str(num_clusters)
    )
    self.cluster_centers = self.add_weight(
      shape=(1, feature_dim, self.num_clusters),
      initializer=tf.keras.initializers.TruncatedNormal(
        stddev=1.0 / math.sqrt(feature_dim)
      ),
      trainable=True,
      name="cluster_centers" + str(num_clusters)
    )
    self.feature_dim = feature_dim
    self.max_frames = input_shape[-2]

  def call(self, frames):
    """Apply the NetVLAD module to the given frames.
    
    Args:
      frames: A tensor with shape [batch_size, max_frames, feature_dim].
    Returns:
      vlad_out: A tensor with shape [batch_size, feature_dim * num_clusters].
    Raises:
      ValueError: If the `feature_dim` of input is not defined.
    """
    feature_dim = self.feature_dim
    max_frames = self.max_frames
    
    frames = tf.reshape(frames, (-1, feature_dim))
    activation = self.fc(frames)
    activation = tf.reshape(activation, (-1, max_frames, self.num_clusters))
    
    activation_sum = tf.math.reduce_sum(activation, axis=-2, keepdims=True)
    cluster_activation = activation_sum * self.cluster_centers

    frames = tf.reshape(frames, (-1, max_frames, feature_dim))
    activation = tf.transpose(
      tf.matmul(tf.transpose(activation, perm=(0, 2, 1)), frames), perm=(0, 2, 1)
    )
    
    vlad_out = activation - cluster_activation
    vlad_out = tf.nn.l2_normalize(vlad_out, 1)
    vlad_out = tf.reshape(vlad_out, (-1, feature_dim * self.num_clusters))
    vlad_out = tf.nn.l2_normalize(vlad_out, 1)

    return vlad_out

  def compute_output_shape(self, input_shape):
    input_shape = tf.TensorShape(input_shape).as_list()
    return tf.TensorShape([input_shape[0], input_shape[-1] * self.num_clusters])

  def get_config(self):
    config = {"num_clusters": self.num_clusters}
    base_config = super().get_config()
    return dict(list(base_config.items()) + list(config.items()))

class ContextGating(tf.keras.layers.Layer):
  """Implements the Context Gating Layer from https://arxiv.org/abs/1706.06905

  Input shape:
  2D tensor with shape: `(batch_size, feature_dim)`.
  Output shape:
  2D tensor with shape: `(batch_size, feature_dim)`.
  """

  def __init__(self, input_shape, **kwargs):
    super(ContextGating, self).__init__(**kwargs)
    feature_dim = input_shape[-1]
    if not isinstance(feature_dim, int):
      feature_dim = feature_dim.value
    self.fc = tf.keras.layers.Dense(
      units=feature_dim,
      activation=tf.nn.sigmoid,
      kernel_regularizer=tf.keras.regularizers.l2(1e-5),
    )

  def call(self, model_input):
    """Apply the ContextGating module to the given input.

    Args:
      model_input: A tensor with shape [batch_size, feature_dim].
    Returns:
      A tensor with shape [batch_size, feature_dim].
    Raises:
      ValueError: If the `feature_dim` of model_input is not defined.
    """
    model_input.shape.assert_has_rank(2)
    feature_dim = model_input.shape.as_list()[-1]
    if feature_dim is None:
      raise ValueError("Last dimension must be defined.")
    context_gate = self.fc(model_input)
    output = tf.math.multiply(context_gate, model_input)
    return output

  def compute_output_shape(self, input_shape):
    return input_shape

  def get_config(self):
    base_config = super().get_config()
    return dict(list(base_config.items()))

class MOELogistic(tf.keras.layers.Layer):
  """Implements a Mixture of Logistic Experts classifier.

  Input shape:
    2D tensor with shape: `(batch_size, feature_dim)`.
  Output shape:
    2D tensor with shape: `(batch_size, num_classes)`.
  """

  def __init__(self, input_shape, num_classes, num_mixtures, **kwargs):
    super(MOELogistic, self).__init__(**kwargs)
    self.num_classes = num_classes
    self.num_mixtures = num_mixtures

    self.gate_fc = tf.keras.layers.Dense(
      units=num_classes*(num_mixtures+1),
      kernel_regularizer=tf.keras.regularizers.l2(1e-6),
    )

    self.expert_fc = tf.keras.layers.Dense(
      units=num_classes*num_mixtures,
      kernel_regularizer=tf.keras.regularizers.l2(1e-6),
    )

  def call(self, input):
    """Apply the MoE algorithm to the given input.

    Args:
      input: A tensor with shape [batch_size, feature_dim].
    Returns:
      A tensor with shape [batch_size, num_classes].
    Raises:
      ValueError: If the `feature_dim` of input is not defined.
    """
    gate_activations = self.gate_fc(input)
    expert_activations = self.expert_fc(input)

    #Calculate the distribution across mixtures
    gate_dist = tf.nn.softmax(tf.reshape(gate_activations, [-1, self.num_mixtures+1]))
    expert_dist = tf.nn.sigmoid(tf.reshape(expert_activations, [-1, self.num_mixtures]))

    probs = tf.reduce_sum(tf.math.multiply(gate_dist[:,:self.num_mixtures], expert_dist),1)
    probs = tf.reshape(probs, [-1, self.num_classes])

    return probs

  def compute_output_shape(self, input_shape):
    return (input_shape[0], self.num_classes)

  def get_config(self):
    base_config = super().get_config()
    config = base_config.update({'number of classes': self.num_classes, 'number of mixtures': self.num_mixtures})
    return config

class VideoClassifier:
  """The Video Classifier model, implemented according to the winning model from the Youtube-8M Challenge.
  The model can be found here: https://arxiv.org/pdf/1706.06905.pdf
  
  Arguments:
    num_clusters: the number of clusters to be used for NetVLAD. The audio clusters will be num_clusters/2.
    video_input_shape: shape of the input video features. Shape of [batch_size, num_samples, video_feature_dim].
    audio_input_shape: shape fo the input audio features. Shape of [batch_size, num_samples, audio_feature_dim].
  
  Raises:
    ValueError: If num_clusters is not divisible by 2.
    ValueError: If the batch sizes of the audio_input_shape and video_input_shape do not match.
    ValueError: If the number of samples of the audio_input_shape and video_input_shape do not match.
  """
  def __init__(self, num_clusters, video_input_shape, audio_input_shape, num_classes, num_mixtures, fc_units, iterations, **kwargs):
    super(VideoClassifier, self).__init__(**kwargs)
    if num_clusters % 2 != 0:
      raise ValueError("num_clusters must be divisible by 2.")
    batch_size = video_input_shape[0]
    if audio_input_shape[0] != batch_size:
      raise ValueError("audio_input_shape[0] must equal video_input_shape[0]. Batch sizes must equal.")
    if audio_input_shape[1] != video_input_shape[1]:
      raise ValueError("audio_input_shape[1] must equal video_input_shape[1]. Number of samples must equal.")

    self.num_frames = video_input_shape[1]
    self.num_classes = num_classes
    self.num_mixtures = num_mixtures
    self.iterations = iterations

    self.video_feature_dim = video_input_shape[2]

    self.video_vlad = NetVLAD(num_clusters, input_shape=video_input_shape, name="video_vlad")
    self.audio_vlad = NetVLAD(num_clusters//2, input_shape=audio_input_shape, name="audio_vlad")

    #Relu6 is used as it is employed in the paper.
    self.fc = tf.keras.layers.Dense(
      units=fc_units,
      activation=tf.nn.relu6,
      kernel_regularizer=tf.keras.regularizers.l2(1e-5),
      name="main_fc"
    )

    self.first_cg = ContextGating(input_shape=(batch_size, fc_units), name="first_cg")

    self.moe = MOELogistic(input_shape=self.first_cg.compute_output_shape((batch_size, fc_units)), num_classes=self.num_classes, num_mixtures=self.num_mixtures, name="moe")

    self.second_cg = ContextGating(input_shape=self.moe.compute_output_shape((batch_size, fc_units)), name="second_cg")

  def build_model(self, input_shape, batch_size):
    """Perform one forward pass of the model.

    Args:
      model_input: input features of shape [batch_size, max_frames, video_feature_dim + audio_feature_dim].
    Returns:
      A tensor with shape [batch_size, num_classes].
    """
    model_input = tf.keras.layers.Input(shape=input_shape, batch_size=batch_size)

    video_input = model_input[:,:,:self.video_feature_dim]
    audio_input = model_input[:,:,self.video_feature_dim:]

    video_vlad_out = self.video_vlad(video_input)
    audio_vlad_out = self.audio_vlad(audio_input)

    vlad_out = tf.concat([video_vlad_out, audio_vlad_out], axis=1)

    fc_out = self.fc(vlad_out)
    cg_out = self.first_cg(fc_out)
    moe_out = self.moe(cg_out)
    final_out = self.second_cg(moe_out)
    
    final_model = tf.keras.models.Model(inputs=model_input, outputs=final_out)

    return final_model

class SegmentClassifier:
  """The Segment Classifier model, implemented according to the winning model from the Youtube-8M Challenge.
  The model can be found here: https://arxiv.org/abs/1911.08548
  
  Arguments:
    num_clusters: the number of clusters to be used for NetVLAD. The audio clusters will be num_clusters/2.
    video_input_shape: shape of the input video features. Shape of [batch_size, num_samples, video_feature_dim].
    audio_input_shape: shape fo the input audio features. Shape of [batch_size, num_samples, audio_feature_dim].
  
  Raises:
    ValueError: If num_clusters is not divisible by 2.
    ValueError: If the batch sizes of the audio_input_shape and video_input_shape do not match.
    ValueError: If the number of samples of the audio_input_shape and video_input_shape do not match.
  """
  def __init__(self, num_clusters, video_input_shape, audio_input_shape, num_classes, fc_units, **kwargs):
    super(SegmentClassifier, self).__init__(**kwargs)
    if num_clusters % 2 != 0:
      raise ValueError("num_clusters must be divisible by 2.")
    batch_size = video_input_shape[0]
    if audio_input_shape[0] != batch_size:
      raise ValueError("audio_input_shape[0] must equal video_input_shape[0]. Batch sizes must equal.")
    if audio_input_shape[1] != video_input_shape[1]:
      raise ValueError("audio_input_shape[1] must equal video_input_shape[1]. Number of samples must equal.")

    self.num_frames = video_input_shape[1]
    self.num_classes = num_classes

    self.video_feature_dim = video_input_shape[2]

    self.video_vlad = NetVLAD(num_clusters, input_shape=video_input_shape, name="video_vlad")
    self.audio_vlad = NetVLAD(num_clusters//2, input_shape=audio_input_shape, name="audio_vlad")

    #Relu6 is used as it is employed in the paper.
    self.fc = tf.keras.layers.Dense(
      units=fc_units,
      activation=tf.nn.relu6,
      kernel_regularizer=tf.keras.regularizers.l2(1e-5),
      name="main_fc"
    )

    self.first_cg = ContextGating(input_shape=(batch_size, fc_units), name="first_cg")

    self.fc2 = tf.keras.layers.Dense(
      units=1,
      activation=tf.keras.activations.sigmoid,
      kernel_regularizer=tf.keras.regularizers.l2(1e-5),
      name="main_fc2"
    )

  def build_model(self, input_shape, second_input_shape, batch_size):
    """Perform one forward pass of the model.

    Args:
      input_shape: input shape for video features. Shape is of the form: [max_frames, video_feature_dim + audio_feature_dim].
      second_input_shape: input shape of new class specific features. Shape is of the form [num_new_features]
    Returns:
      A tensor with shape [batch_size, num_classes].
    """
    model_input = tf.keras.layers.Input(shape=input_shape, batch_size=batch_size)
    model_input2 = tf.keras.layers.Input(shape=second_input_shape, batch_size=batch_size)
    video_input = model_input[:,:,:self.video_feature_dim]
    audio_input = model_input[:,:,self.video_feature_dim:]
    video_vlad_out = self.video_vlad(video_input)
    audio_vlad_out = self.audio_vlad(audio_input)
    vlad_out = tf.concat([video_vlad_out, audio_vlad_out], axis=1)
    vlad_out = tf.concat([vlad_out, model_input2], axis=1)
    fc_out = self.fc(vlad_out)
    cg_out = self.first_cg(fc_out)
    final_out = self.fc2(cg_out)
    final_model = tf.keras.models.Model(inputs=[model_input, model_input2], outputs=final_out)
    return final_model
[(68, 'tensorflow.reshape', 'tf.reshape', True, 'import tensorflow as tf\n'), (70, 'tensorflow.reshape', 'tf.reshape', True, 'import tensorflow as tf\n'), (72, 'tensorflow.math.reduce_sum', 'tf.math.reduce_sum', True, 'import tensorflow as tf\n'), (75, 'tensorflow.reshape', 'tf.reshape', True, 'import tensorflow as tf\n'), (81, 'tensorflow.nn.l2_normalize', 'tf.nn.l2_normalize', True, 'import tensorflow as tf\n'), (82, 'tensorflow.reshape', 'tf.reshape', True, 'import tensorflow as tf\n'), (83, 'tensorflow.nn.l2_normalize', 'tf.nn.l2_normalize', True, 'import tensorflow as tf\n'), (89, 'tensorflow.TensorShape', 'tf.TensorShape', True, 'import tensorflow as tf\n'), (131, 'tensorflow.math.multiply', 'tf.math.multiply', True, 'import tensorflow as tf\n'), (183, 'tensorflow.reshape', 'tf.reshape', True, 'import tensorflow as tf\n'), (251, 'tensorflow.keras.layers.Input', 'tf.keras.layers.Input', True, 'import tensorflow as tf\n'), (259, 'tensorflow.concat', 'tf.concat', True, 'import tensorflow as tf\n'), (266, 'tensorflow.keras.models.Model', 'tf.keras.models.Model', True, 'import tensorflow as tf\n'), (328, 'tensorflow.keras.layers.Input', 'tf.keras.layers.Input', True, 'import tensorflow as tf\n'), (329, 'tensorflow.keras.layers.Input', 'tf.keras.layers.Input', True, 'import tensorflow as tf\n'), (334, 'tensorflow.concat', 'tf.concat', True, 'import tensorflow as tf\n'), (335, 'tensorflow.concat', 'tf.concat', True, 'import tensorflow as tf\n'), (339, 'tensorflow.keras.models.Model', 'tf.keras.models.Model', True, 'import tensorflow as tf\n'), (179, 'tensorflow.reshape', 'tf.reshape', True, 'import tensorflow as tf\n'), (180, 'tensorflow.reshape', 'tf.reshape', True, 'import tensorflow as tf\n'), (182, 'tensorflow.math.multiply', 'tf.math.multiply', True, 'import tensorflow as tf\n'), (41, 'tensorflow.keras.regularizers.l2', 'tf.keras.regularizers.l2', True, 'import tensorflow as tf\n'), (77, 'tensorflow.transpose', 'tf.transpose', True, 'import tensorflow as tf\n'), (88, 'tensorflow.TensorShape', 'tf.TensorShape', True, 'import tensorflow as tf\n'), (113, 'tensorflow.keras.regularizers.l2', 'tf.keras.regularizers.l2', True, 'import tensorflow as tf\n'), (157, 'tensorflow.keras.regularizers.l2', 'tf.keras.regularizers.l2', True, 'import tensorflow as tf\n'), (162, 'tensorflow.keras.regularizers.l2', 'tf.keras.regularizers.l2', True, 'import tensorflow as tf\n'), (233, 'tensorflow.keras.regularizers.l2', 'tf.keras.regularizers.l2', True, 'import tensorflow as tf\n'), (306, 'tensorflow.keras.regularizers.l2', 'tf.keras.regularizers.l2', True, 'import tensorflow as tf\n'), (315, 'tensorflow.keras.regularizers.l2', 'tf.keras.regularizers.l2', True, 'import tensorflow as tf\n'), (47, 'math.sqrt', 'math.sqrt', False, 'import math\n')]