# Copyright 2023 The TensorFlow Authors. 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.

"""Video classification task definition."""
from typing import Dict, List, Optional, Tuple

from absl import logging
import tensorflow as tf

from official.core import base_task
from official.core import input_reader
from official.core import task_factory
from official.modeling import tf_utils
from official.projects.yt8m.configs import yt8m as yt8m_cfg
from official.projects.yt8m.dataloaders import yt8m_input
from official.projects.yt8m.eval_utils import eval_util
from official.projects.yt8m.modeling import yt8m_model_utils as utils
from official.projects.yt8m.modeling.yt8m_model import DbofModel


@task_factory.register_task_cls(yt8m_cfg.YT8MTask)
class YT8MTask(base_task.Task):
  """A task for video classification."""

  def build_model(self):
    """Builds model for YT8M Task."""
    train_cfg = self.task_config.train_data
    common_input_shape = [None, sum(train_cfg.feature_sizes)]

    # [batch_size x num_frames x num_features]
    input_specs = tf.keras.layers.InputSpec(shape=[None] + common_input_shape)
    logging.info('Build model input %r', common_input_shape)

    l2_weight_decay = self.task_config.losses.l2_weight_decay
    # Model configuration.
    model_config = self.task_config.model
    model = DbofModel(
        params=model_config,
        input_specs=input_specs,
        num_classes=train_cfg.num_classes,
        l2_weight_decay=l2_weight_decay)

    non_trainable_batch_norm_variables = []
    non_trainable_extra_variables = []
    for var in model.variables:
      if not var.trainable:
        if 'moving_mean' or 'moving_variance' in var.name:
          non_trainable_batch_norm_variables.append(var)
        else:
          non_trainable_extra_variables.append(var)

    logging.info(
        'Trainable model variables:\n%s',
        '\n'.join(
            [f'{var.name}\t{var.shape}' for var in model.trainable_variables]
        ),
    )
    logging.info(
        (
            'Non-trainable batch norm variables (get updated in training'
            ' mode):\n%s'
        ),
        '\n'.join(
            [
                f'{var.name}\t{var.shape}'
                for var in non_trainable_batch_norm_variables
            ]
        ),
    )
    logging.info(
        'Non-trainable frozen model variables:\n%s',
        '\n'.join(
            [
                f'{var.name}\t{var.shape}'
                for var in non_trainable_extra_variables
            ]
        ),
    )
    return model

  def build_inputs(self, params: yt8m_cfg.DataConfig, input_context=None):
    """Builds input.

    Args:
      params: configuration for input data
      input_context: indicates information about the compute replicas and input
        pipelines

    Returns:
      dataset: dataset fetched from reader
    """

    decoder = yt8m_input.Decoder(input_params=params)
    decoder_fn = decoder.decode
    parser = yt8m_input.Parser(input_params=params)
    parser_fn = parser.parse_fn(params.is_training)
    postprocess = yt8m_input.PostBatchProcessor(input_params=params)
    postprocess_fn = postprocess.post_fn
    transform_batch = yt8m_input.TransformBatcher(input_params=params)
    batch_fn = transform_batch.batch_fn

    reader = input_reader.InputReader(
        params,
        dataset_fn=tf.data.TFRecordDataset,
        decoder_fn=decoder_fn,
        parser_fn=parser_fn,
        postprocess_fn=postprocess_fn,
        transform_and_batch_fn=batch_fn)

    dataset = reader.read(input_context=input_context)

    return dataset

  def build_losses(self,
                   labels,
                   model_outputs,
                   label_weights=None,
                   aux_losses=None):
    """Sigmoid Cross Entropy.

    Args:
      labels: tensor containing truth labels.
      model_outputs: output logits of the classifier.
      label_weights: optional tensor of label weights.
      aux_losses: tensor containing auxiliarly loss tensors, i.e. `losses` in
        keras.Model.

    Returns:
      A dict of tensors contains total loss, model loss tensors.
    """
    losses_config = self.task_config.losses
    model_loss = tf.keras.losses.binary_crossentropy(
        labels,
        model_outputs,
        from_logits=losses_config.from_logits,
        label_smoothing=losses_config.label_smoothing,
        axis=None)

    if label_weights is None:
      model_loss = tf_utils.safe_mean(model_loss)
    else:
      model_loss = model_loss * label_weights
      # Manutally compute weighted mean loss.
      total_loss = tf.reduce_sum(model_loss)
      total_weight = tf.cast(
          tf.reduce_sum(label_weights), dtype=total_loss.dtype)
      model_loss = tf.math.divide_no_nan(total_loss, total_weight)

    total_loss = model_loss
    if aux_losses:
      total_loss += tf.add_n(aux_losses)

    return {'total_loss': total_loss, 'model_loss': model_loss}

  def build_metrics(self, training=True):
    """Gets streaming metrics for training/validation.

       metric: mAP/gAP
       top_k: A positive integer specifying how many predictions are considered
        per video.
       top_n: A positive Integer specifying the average precision at n, or None
        to use all provided data points.
    Args:
      training: Bool value, true for training mode, false for eval/validation.

    Returns:
      A list of metrics to be used.
    """
    metrics = []
    metric_names = ['total_loss', 'model_loss']
    for name in metric_names:
      metrics.append(tf.keras.metrics.Mean(name, dtype=tf.float32))

    if (
        self.task_config.evaluation.average_precision is not None
        and not training
    ):
      # Cannot run in train step.
      num_classes = self.task_config.validation_data.num_classes
      top_k = self.task_config.evaluation.average_precision.top_k
      top_n = self.task_config.evaluation.average_precision.top_n
      self.avg_prec_metric = eval_util.EvaluationMetrics(
          num_classes, top_k=top_k, top_n=top_n)

    return metrics

  def process_metrics(
      self,
      metrics: List[tf.keras.metrics.Metric],
      labels: tf.Tensor,
      outputs: tf.Tensor,
      model_losses: Optional[Dict[str, tf.Tensor]] = None,
      label_weights: Optional[tf.Tensor] = None,
      training: bool = True,
      **kwargs,
  ) -> Dict[str, Tuple[tf.Tensor, ...]]:
    """Updates metrics.

    Args:
      metrics: Evaluation metrics to be updated.
      labels: A tensor containing truth labels.
      outputs: Model output logits of the classifier.
      model_losses: An optional dict of model losses.
      label_weights: Optional label weights, can be broadcast into shape of
        outputs/labels.
      training: Bool indicates if in training mode.
      **kwargs: Additional input arguments.

    Returns:
      Updated dict of metrics log.
    """
    if model_losses is None:
      model_losses = {}

    logs = {}
    if (
        self.task_config.evaluation.average_precision is not None
        and not training
    ):
      logs.update({self.avg_prec_metric.name: (labels, outputs)})

    for m in metrics:
      if m.name in model_losses:
        m.update_state(model_losses[m.name])
        logs[m.name] = m.result()
    return logs

  def _preprocess_model_inputs(self,
                               inputs: dict[str, tf.Tensor],
                               training: bool = True):
    """Preprocesses input tensors before model on device."""
    if training:
      data_config = self.task_config.train_data
    else:
      data_config = self.task_config.validation_data

    features = inputs['video_matrix']
    num_frames = inputs['num_frames']

    # sample random frames / random sequence.
    num_frames = tf.cast(num_frames, tf.float32)
    num_sample_frames = data_config.num_sample_frames
    if self.task_config.model.sample_random_frames:
      features = utils.sample_random_frames(
          features, num_frames, num_sample_frames)
    else:
      features = utils.sample_random_sequence(
          features, num_frames, num_sample_frames)
    return features

  def _preprocess_labels(self,
                         inputs: dict[str, tf.Tensor],
                         training: bool = True):
    """Preprocesses labels."""
    del training  # training is unused in _preprocess_labels in YT8M.
    labels = inputs['labels']
    label_weights = inputs.get('label_weights', None)

    return labels, label_weights

  def _postprocess_outputs(self,
                           inputs,
                           outputs,
                           labels,
                           label_weights,
                           training: bool = True):
    """Postprocess model outputs (inputs / labels / label_weights)."""
    if not training and self.task_config.validation_data.segment_labels:
      # workaround to ignore the unrated labels.
      outputs *= label_weights
      # remove padding
      outputs = outputs[~tf.reduce_all(labels == -1, axis=1)]
      labels = labels[~tf.reduce_all(labels == -1, axis=1)]
    return inputs, outputs, labels, label_weights

  def train_step(self, inputs, model, optimizer, metrics=None):
    """Does forward and backward.

    Args:
      inputs: a dictionary of input tensors. output_dict = { "video_ids":
        batch_video_ids, "video_matrix": batch_video_matrix, "labels":
        batch_labels, "num_frames": batch_frames, }
      model: the model, forward pass definition.
      optimizer: the optimizer for this training step.
      metrics: a nested structure of metrics objects.

    Returns:
      a dictionary of logs.
    """
    model_inputs = self._preprocess_model_inputs(inputs, training=True)
    labels, label_weights = self._preprocess_labels(inputs, training=True)

    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync
    with tf.GradientTape() as tape:
      outputs = model(model_inputs, training=True)
      # Casting output layer as float32 is necessary when mixed_precision is
      # mixed_float16 or mixed_bfloat16 to ensure output is casted as float32.
      outputs = tf.nest.map_structure(lambda x: tf.cast(x, tf.float32), outputs)
      # Post-process model / label outputs.
      inputs, outputs, labels, label_weights = self._postprocess_outputs(
          inputs, outputs, labels, label_weights, training=True)

      # Computes per-replica loss
      all_losses = self.build_losses(
          model_outputs=outputs,
          labels=labels,
          label_weights=label_weights,
          aux_losses=model.losses)

      loss = all_losses['total_loss']
      # Scales loss as the default gradients allreduce performs sum inside the
      # optimizer.
      scaled_loss = loss / num_replicas

      # For mixed_precision policy, when LossScaleOptimizer is used, loss is
      # scaled for numerical stability.
      if isinstance(optimizer, tf.keras.mixed_precision.LossScaleOptimizer):
        scaled_loss = optimizer.get_scaled_loss(scaled_loss)

    tvars = model.trainable_variables
    grads = tape.gradient(scaled_loss, tvars)
    # Scales back gradient before apply_gradients when LossScaleOptimizer is
    # used.
    if isinstance(optimizer, tf.keras.mixed_precision.LossScaleOptimizer):
      grads = optimizer.get_unscaled_gradients(grads)

    # Apply gradient clipping.
    if self.task_config.gradient_clip_norm > 0:
      grads, _ = tf.clip_by_global_norm(grads,
                                        self.task_config.gradient_clip_norm)
    optimizer.apply_gradients(list(zip(grads, tvars)))

    logs = {self.loss: loss}
    logs.update(
        self.process_metrics(
            metrics,
            labels=labels,
            outputs=outputs,
            model_losses=all_losses,
            label_weights=label_weights,
            training=True))
    return logs

  def validation_step(self, inputs, model, metrics=None):
    """Validatation step.

    Args:
      inputs: a dictionary of input tensors. output_dict = { "video_ids":
        batch_video_ids, "video_matrix": batch_video_matrix, "labels":
        batch_labels, "num_frames": batch_frames, }
      model: the model, forward definition
      metrics: a nested structure of metrics objects.

    Returns:
      a dictionary of logs.
    """
    model_inputs = self._preprocess_model_inputs(inputs, training=False)
    labels, label_weights = self._preprocess_labels(inputs, training=False)

    outputs = self.inference_step(model_inputs, model)
    outputs = tf.nest.map_structure(lambda x: tf.cast(x, tf.float32), outputs)
    inputs, outputs, labels, label_weights = self._postprocess_outputs(
        inputs, outputs, labels, label_weights, training=False)

    all_losses = self.build_losses(
        labels=labels,
        model_outputs=outputs,
        label_weights=label_weights,
        aux_losses=model.losses)

    logs = {self.loss: all_losses['total_loss']}
    logs.update(
        self.process_metrics(
            metrics,
            labels=labels,
            outputs=outputs,
            model_losses=all_losses,
            label_weights=inputs.get('label_weights', None),
            training=False))

    return logs

  def inference_step(self, inputs, model):
    """Performs the forward step."""
    return model(inputs, training=False)

  def aggregate_logs(self, state=None, step_logs=None):
    if self.task_config.evaluation.average_precision is not None:
      if state is None:
        state = self.avg_prec_metric
      self.avg_prec_metric.accumulate(
          labels=step_logs[self.avg_prec_metric.name][0],
          predictions=step_logs[self.avg_prec_metric.name][1])
    return state

  def reduce_aggregated_logs(self, aggregated_logs, global_step=None):
    if self.task_config.evaluation.average_precision is not None:
      avg_prec_metrics = self.avg_prec_metric.get()
      self.avg_prec_metric.clear()
      return avg_prec_metrics
    return None
