# 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.

"""Contains definitions of Mobile Video Networks.

Reference: https://arxiv.org/pdf/2103.11511.pdf
"""
import dataclasses
import math
from typing import Dict, Mapping, Optional, Sequence, Tuple, Union

from absl import logging
import tensorflow as tf, tf_keras

from official.modeling import hyperparams
from official.projects.movinet.modeling import movinet_layers
from official.vision.modeling.backbones import factory

# Defines a set of kernel sizes and stride sizes to simplify and shorten
# architecture definitions for configs below.
KernelSize = Tuple[int, int, int]

# K(ab) represents a 3D kernel of size (a, b, b)
K13: KernelSize = (1, 3, 3)
K15: KernelSize = (1, 5, 5)
K33: KernelSize = (3, 3, 3)
K53: KernelSize = (5, 3, 3)

# S(ab) represents a 3D stride of size (a, b, b)
S11: KernelSize = (1, 1, 1)
S12: KernelSize = (1, 2, 2)
S22: KernelSize = (2, 2, 2)
S21: KernelSize = (2, 1, 1)

# Type for a state container (map)
TensorMap = Mapping[str, tf.Tensor]


@dataclasses.dataclass
class BlockSpec:
  """Configuration of a block."""


@dataclasses.dataclass
class StemSpec(BlockSpec):
  """Configuration of a Movinet block."""
  filters: int = 0
  kernel_size: KernelSize = (0, 0, 0)
  strides: KernelSize = (0, 0, 0)


@dataclasses.dataclass
class MovinetBlockSpec(BlockSpec):
  """Configuration of a Movinet block."""
  base_filters: int = 0
  expand_filters: Sequence[int] = ()
  kernel_sizes: Sequence[KernelSize] = ()
  strides: Sequence[KernelSize] = ()


@dataclasses.dataclass
class HeadSpec(BlockSpec):
  """Configuration of a Movinet block."""
  project_filters: int = 0
  head_filters: int = 0


# Block specs specify the architecture of each model
BLOCK_SPECS = {
    'a0': (
        StemSpec(filters=8, kernel_size=K13, strides=S12),
        MovinetBlockSpec(
            base_filters=8,
            expand_filters=(24,),
            kernel_sizes=(K15,),
            strides=(S12,)),
        MovinetBlockSpec(
            base_filters=32,
            expand_filters=(80, 80, 80),
            kernel_sizes=(K33, K33, K33),
            strides=(S12, S11, S11)),
        MovinetBlockSpec(
            base_filters=56,
            expand_filters=(184, 112, 184),
            kernel_sizes=(K53, K33, K33),
            strides=(S12, S11, S11)),
        MovinetBlockSpec(
            base_filters=56,
            expand_filters=(184, 184, 184, 184),
            kernel_sizes=(K53, K33, K33, K33),
            strides=(S11, S11, S11, S11)),
        MovinetBlockSpec(
            base_filters=104,
            expand_filters=(384, 280, 280, 344),
            kernel_sizes=(K53, K15, K15, K15),
            strides=(S12, S11, S11, S11)),
        HeadSpec(project_filters=480, head_filters=2048),
    ),
    'a1': (
        StemSpec(filters=16, kernel_size=K13, strides=S12),
        MovinetBlockSpec(
            base_filters=16,
            expand_filters=(40, 40),
            kernel_sizes=(K15, K33),
            strides=(S12, S11)),
        MovinetBlockSpec(
            base_filters=40,
            expand_filters=(96, 120, 96, 96),
            kernel_sizes=(K33, K33, K33, K33),
            strides=(S12, S11, S11, S11)),
        MovinetBlockSpec(
            base_filters=64,
            expand_filters=(216, 128, 216, 168, 216),
            kernel_sizes=(K53, K33, K33, K33, K33),
            strides=(S12, S11, S11, S11, S11)),
        MovinetBlockSpec(
            base_filters=64,
            expand_filters=(216, 216, 216, 128, 128, 216),
            kernel_sizes=(K53, K33, K33, K33, K15, K33),
            strides=(S11, S11, S11, S11, S11, S11)),
        MovinetBlockSpec(
            base_filters=136,
            expand_filters=(456, 360, 360, 360, 456, 456, 544),
            kernel_sizes=(K53, K15, K15, K15, K15, K33, K13),
            strides=(S12, S11, S11, S11, S11, S11, S11)),
        HeadSpec(project_filters=600, head_filters=2048),
    ),
    'a2': (
        StemSpec(filters=16, kernel_size=K13, strides=S12),
        MovinetBlockSpec(
            base_filters=16,
            expand_filters=(40, 40, 64),
            kernel_sizes=(K15, K33, K33),
            strides=(S12, S11, S11)),
        MovinetBlockSpec(
            base_filters=40,
            expand_filters=(96, 120, 96, 96, 120),
            kernel_sizes=(K33, K33, K33, K33, K33),
            strides=(S12, S11, S11, S11, S11)),
        MovinetBlockSpec(
            base_filters=72,
            expand_filters=(240, 160, 240, 192, 240),
            kernel_sizes=(K53, K33, K33, K33, K33),
            strides=(S12, S11, S11, S11, S11)),
        MovinetBlockSpec(
            base_filters=72,
            expand_filters=(240, 240, 240, 240, 144, 240),
            kernel_sizes=(K53, K33, K33, K33, K15, K33),
            strides=(S11, S11, S11, S11, S11, S11)),
        MovinetBlockSpec(
            base_filters=144,
            expand_filters=(480, 384, 384, 480, 480, 480, 576),
            kernel_sizes=(K53, K15, K15, K15, K15, K33, K13),
            strides=(S12, S11, S11, S11, S11, S11, S11)),
        HeadSpec(project_filters=640, head_filters=2048),
    ),
    'a3': (
        StemSpec(filters=16, kernel_size=K13, strides=S12),
        MovinetBlockSpec(
            base_filters=16,
            expand_filters=(40, 40, 64, 40),
            kernel_sizes=(K15, K33, K33, K33),
            strides=(S12, S11, S11, S11)),
        MovinetBlockSpec(
            base_filters=48,
            expand_filters=(112, 144, 112, 112, 144, 144),
            kernel_sizes=(K33, K33, K33, K15, K33, K33),
            strides=(S12, S11, S11, S11, S11, S11)),
        MovinetBlockSpec(
            base_filters=80,
            expand_filters=(240, 152, 240, 192, 240),
            kernel_sizes=(K53, K33, K33, K33, K33),
            strides=(S12, S11, S11, S11, S11)),
        MovinetBlockSpec(
            base_filters=88,
            expand_filters=(264, 264, 264, 264, 160, 264, 264, 264),
            kernel_sizes=(K53, K33, K33, K33, K15, K33, K33, K33),
            strides=(S11, S11, S11, S11, S11, S11, S11, S11)),
        MovinetBlockSpec(
            base_filters=168,
            expand_filters=(560, 448, 448, 560, 560, 560, 448, 448, 560, 672),
            kernel_sizes=(K53, K15, K15, K15, K15, K33, K15, K15, K33, K13),
            strides=(S12, S11, S11, S11, S11, S11, S11, S11, S11, S11)),
        HeadSpec(project_filters=744, head_filters=2048),
    ),
    'a4': (
        StemSpec(filters=24, kernel_size=K13, strides=S12),
        MovinetBlockSpec(
            base_filters=24,
            expand_filters=(64, 64, 96, 64, 96, 64),
            kernel_sizes=(K15, K33, K33, K33, K33, K33),
            strides=(S12, S11, S11, S11, S11, S11)),
        MovinetBlockSpec(
            base_filters=56,
            expand_filters=(168, 168, 136, 136, 168, 168, 168, 136, 136),
            kernel_sizes=(K33, K33, K33, K33, K33, K33, K33, K15, K33),
            strides=(S12, S11, S11, S11, S11, S11, S11, S11, S11)),
        MovinetBlockSpec(
            base_filters=96,
            expand_filters=(320, 160, 320, 192, 320, 160, 320, 256, 320),
            kernel_sizes=(K53, K33, K33, K33, K33, K33, K33, K33, K33),
            strides=(S12, S11, S11, S11, S11, S11, S11, S11, S11)),
        MovinetBlockSpec(
            base_filters=96,
            expand_filters=(320, 320, 320, 320, 192, 320, 320, 192, 320, 320),
            kernel_sizes=(K53, K33, K33, K33, K15, K33, K33, K33, K33, K33),
            strides=(S11, S11, S11, S11, S11, S11, S11, S11, S11, S11)),
        MovinetBlockSpec(
            base_filters=192,
            expand_filters=(640, 512, 512, 640, 640, 640, 512, 512, 640, 768,
                            640, 640, 768),
            kernel_sizes=(K53, K15, K15, K15, K15, K33, K15, K15, K15, K15, K15,
                          K33, K33),
            strides=(S12, S11, S11, S11, S11, S11, S11, S11, S11, S11, S11, S11,
                     S11)),
        HeadSpec(project_filters=856, head_filters=2048),
    ),
    'a5': (
        StemSpec(filters=24, kernel_size=K13, strides=S12),
        MovinetBlockSpec(
            base_filters=24,
            expand_filters=(64, 64, 96, 64, 96, 64),
            kernel_sizes=(K15, K15, K33, K33, K33, K33),
            strides=(S12, S11, S11, S11, S11, S11)),
        MovinetBlockSpec(
            base_filters=64,
            expand_filters=(192, 152, 152, 152, 192, 192, 192, 152, 152, 192,
                            192),
            kernel_sizes=(K53, K33, K33, K33, K33, K33, K33, K33, K33, K33,
                          K33),
            strides=(S12, S11, S11, S11, S11, S11, S11, S11, S11, S11, S11)),
        MovinetBlockSpec(
            base_filters=112,
            expand_filters=(376, 224, 376, 376, 296, 376, 224, 376, 376, 296,
                            376, 376, 376),
            kernel_sizes=(K53, K33, K33, K33, K33, K33, K33, K33, K33, K33, K33,
                          K33, K33),
            strides=(S12, S11, S11, S11, S11, S11, S11, S11, S11, S11, S11, S11,
                     S11)),
        MovinetBlockSpec(
            base_filters=120,
            expand_filters=(376, 376, 376, 376, 224, 376, 376, 224, 376, 376,
                            376),
            kernel_sizes=(K53, K33, K33, K33, K15, K33, K33, K33, K33, K33,
                          K33),
            strides=(S11, S11, S11, S11, S11, S11, S11, S11, S11, S11, S11)),
        MovinetBlockSpec(
            base_filters=224,
            expand_filters=(744, 744, 600, 600, 744, 744, 744, 896, 600, 600,
                            896, 744, 744, 896, 600, 600, 744, 744),
            kernel_sizes=(K53, K33, K15, K15, K15, K15, K33, K15, K15, K15, K15,
                          K15, K33, K15, K15, K15, K15, K33),
            strides=(S12, S11, S11, S11, S11, S11, S11, S11, S11, S11, S11, S11,
                     S11, S11, S11, S11, S11, S11)),
        HeadSpec(project_filters=992, head_filters=2048),
    ),
    't0': (
        StemSpec(filters=8, kernel_size=K13, strides=S12),
        MovinetBlockSpec(
            base_filters=8,
            expand_filters=(16,),
            kernel_sizes=(K15,),
            strides=(S12,)),
        MovinetBlockSpec(
            base_filters=32,
            expand_filters=(72, 72),
            kernel_sizes=(K33, K15),
            strides=(S12, S11)),
        MovinetBlockSpec(
            base_filters=56,
            expand_filters=(112, 112, 112),
            kernel_sizes=(K53, K15, K33),
            strides=(S12, S11, S11)),
        MovinetBlockSpec(
            base_filters=56,
            expand_filters=(184, 184, 184, 184),
            kernel_sizes=(K53, K15, K33, K33),
            strides=(S11, S11, S11, S11)),
        MovinetBlockSpec(
            base_filters=104,
            expand_filters=(344, 344, 344, 344),
            kernel_sizes=(K53, K15, K15, K33),
            strides=(S12, S11, S11, S11)),
        HeadSpec(project_filters=240, head_filters=1024),
    ),
}


@tf_keras.utils.register_keras_serializable(package='Vision')
class Movinet(tf_keras.Model):
  """Class to build Movinet family model.

  Reference: https://arxiv.org/pdf/2103.11511.pdf
  """

  def __init__(self,
               model_id: str = 'a0',
               causal: bool = False,
               use_positional_encoding: bool = False,
               conv_type: str = '3d',
               se_type: str = '3d',
               input_specs: Optional[tf_keras.layers.InputSpec] = None,
               activation: str = 'swish',
               gating_activation: str = 'sigmoid',
               use_sync_bn: bool = True,
               norm_momentum: float = 0.99,
               norm_epsilon: float = 0.001,
               kernel_initializer: str = 'HeNormal',
               kernel_regularizer: Optional[str] = None,
               bias_regularizer: Optional[str] = None,
               stochastic_depth_drop_rate: float = 0.,
               use_external_states: bool = False,
               output_states: bool = True,
               average_pooling_type: str = '3d',
               **kwargs):
    """MoViNet initialization function.

    Args:
      model_id: name of MoViNet backbone model.
      causal: use causal mode, with CausalConv and CausalSE operations.
      use_positional_encoding:  if True, adds a positional encoding before
          temporal convolutions and the cumulative global average pooling
          layers.
      conv_type: '3d', '2plus1d', or '3d_2plus1d'. '3d' configures the network
        to use the default 3D convolution. '2plus1d' uses (2+1)D convolution
        with Conv2D operations and 2D reshaping (e.g., a 5x3x3 kernel becomes
        3x3 followed by 5x1 conv). '3d_2plus1d' uses (2+1)D convolution with
        Conv3D and no 2D reshaping (e.g., a 5x3x3 kernel becomes 1x3x3 followed
        by 5x1x1 conv).
      se_type: '3d', '2d', '2plus3d' or 'none'. '3d' uses the default 3D
          spatiotemporal global average pooling for squeeze excitation. '2d'
          uses 2D spatial global average pooling  on each frame. '2plus3d'
          concatenates both 3D and 2D global average pooling.
      input_specs: the model input spec to use.
      activation: name of the main activation function.
      gating_activation: gating activation to use in squeeze excitation layers.
      use_sync_bn: if True, use synchronized batch normalization.
      norm_momentum: normalization momentum for the moving average.
      norm_epsilon: small float added to variance to avoid dividing by
        zero.
      kernel_initializer: kernel_initializer for convolutional layers.
      kernel_regularizer: tf_keras.regularizers.Regularizer object for Conv2D.
        Defaults to None.
      bias_regularizer: tf_keras.regularizers.Regularizer object for Conv2d.
        Defaults to None.
      stochastic_depth_drop_rate: the base rate for stochastic depth.
      use_external_states: if True, expects states to be passed as additional
        input.
      output_states: if True, output intermediate states that can be used to run
          the model in streaming mode. Inputting the output states of the
          previous input clip with the current input clip will utilize a stream
          buffer for streaming video.
      average_pooling_type: The average pooling type. Currently supporting
        ['3d', '2d', 'none'].
      **kwargs: keyword arguments to be passed.
    """
    block_specs = BLOCK_SPECS[model_id]
    if input_specs is None:
      input_specs = tf_keras.layers.InputSpec(shape=[None, None, None, None, 3])

    if conv_type not in ('3d', '2plus1d', '3d_2plus1d'):
      raise ValueError('Unknown conv type: {}'.format(conv_type))
    if se_type not in ('3d', '2d', '2plus3d', 'none'):
      raise ValueError('Unknown squeeze excitation type: {}'.format(se_type))

    self._model_id = model_id
    self._block_specs = block_specs
    self._causal = causal
    self._use_positional_encoding = use_positional_encoding
    self._conv_type = conv_type
    self._se_type = se_type
    self._input_specs = input_specs
    self._use_sync_bn = use_sync_bn
    self._activation = activation
    self._gating_activation = gating_activation
    self._norm_momentum = norm_momentum
    self._norm_epsilon = norm_epsilon
    self._norm = tf_keras.layers.BatchNormalization
    self._kernel_initializer = kernel_initializer
    self._kernel_regularizer = kernel_regularizer
    self._bias_regularizer = bias_regularizer
    self._stochastic_depth_drop_rate = stochastic_depth_drop_rate
    self._use_external_states = use_external_states
    self._output_states = output_states
    self._average_pooling_type = average_pooling_type

    if self._use_external_states and not self._causal:
      raise ValueError('External states should be used with causal mode.')
    if not isinstance(block_specs[0], StemSpec):
      raise ValueError(
          'Expected first spec to be StemSpec, got {}'.format(block_specs[0]))
    if not isinstance(block_specs[-1], HeadSpec):
      raise ValueError(
          'Expected final spec to be HeadSpec, got {}'.format(block_specs[-1]))
    self._head_filters = block_specs[-1].head_filters

    state_specs = None
    if use_external_states:
      self._set_dtype_policy(input_specs.dtype)
      state_specs = self.initial_state_specs(input_specs.shape)

    inputs, outputs = self._build_network(input_specs, state_specs=state_specs)

    super(Movinet, self).__init__(inputs=inputs, outputs=outputs, **kwargs)

    self._state_specs = state_specs

  def _build_network(
      self,
      input_specs: tf_keras.layers.InputSpec,
      state_specs: Optional[Mapping[str, tf_keras.layers.InputSpec]] = None,
  ) -> Tuple[TensorMap, Union[TensorMap, Tuple[TensorMap, TensorMap]]]:
    """Builds the model network.

    Args:
      input_specs: the model input spec to use.
      state_specs: a dict mapping a state name to the corresponding state spec.
        State names should match with the `state` input/output dict.

    Returns:
      Inputs and outputs as a tuple. Inputs are expected to be a dict with
      base input and states. Outputs are expected to be a dict of endpoints
      and (optional) output states.
    """
    state_specs = state_specs if state_specs is not None else {}

    image_input = tf_keras.Input(shape=input_specs.shape[1:], name='inputs')

    states = {
        name: tf_keras.Input(shape=spec.shape[1:], dtype=spec.dtype, name=name)
        for name, spec in state_specs.items()
    }

    inputs = {**states, 'image': image_input}
    endpoints = {}

    x = image_input

    num_layers = sum(
        len(block.expand_filters)
        for block in self._block_specs
        if isinstance(block, MovinetBlockSpec))
    stochastic_depth_idx = 1
    for block_idx, block in enumerate(self._block_specs):
      if isinstance(block, StemSpec):
        layer_obj = movinet_layers.Stem(
            block.filters,
            block.kernel_size,
            block.strides,
            conv_type=self._conv_type,
            causal=self._causal,
            activation=self._activation,
            kernel_initializer=self._kernel_initializer,
            kernel_regularizer=self._kernel_regularizer,
            batch_norm_layer=self._norm,
            batch_norm_momentum=self._norm_momentum,
            batch_norm_epsilon=self._norm_epsilon,
            use_sync_bn=self._use_sync_bn,
            state_prefix='state_stem',
            name='stem')
        x, states = layer_obj(x, states=states)
        endpoints['stem'] = x
      elif isinstance(block, MovinetBlockSpec):
        if not (len(block.expand_filters) == len(block.kernel_sizes) ==
                len(block.strides)):
          raise ValueError(
              'Lengths of block parameters differ: {}, {}, {}'.format(
                  len(block.expand_filters),
                  len(block.kernel_sizes),
                  len(block.strides)))
        params = list(zip(block.expand_filters,
                          block.kernel_sizes,
                          block.strides))
        for layer_idx, layer in enumerate(params):
          stochastic_depth_drop_rate = (
              self._stochastic_depth_drop_rate * stochastic_depth_idx /
              num_layers)
          expand_filters, kernel_size, strides = layer
          name = f'block{block_idx-1}_layer{layer_idx}'
          layer_obj = movinet_layers.MovinetBlock(
              block.base_filters,
              expand_filters,
              kernel_size=kernel_size,
              strides=strides,
              causal=self._causal,
              activation=self._activation,
              gating_activation=self._gating_activation,
              stochastic_depth_drop_rate=stochastic_depth_drop_rate,
              conv_type=self._conv_type,
              se_type=self._se_type,
              use_positional_encoding=
              self._use_positional_encoding and self._causal,
              kernel_initializer=self._kernel_initializer,
              kernel_regularizer=self._kernel_regularizer,
              batch_norm_layer=self._norm,
              batch_norm_momentum=self._norm_momentum,
              batch_norm_epsilon=self._norm_epsilon,
              use_sync_bn=self._use_sync_bn,
              state_prefix=f'state_{name}',
              name=name)
          x, states = layer_obj(x, states=states)

          endpoints[name] = x
          stochastic_depth_idx += 1
      elif isinstance(block, HeadSpec):
        layer_obj = movinet_layers.Head(
            project_filters=block.project_filters,
            conv_type=self._conv_type,
            activation=self._activation,
            kernel_initializer=self._kernel_initializer,
            kernel_regularizer=self._kernel_regularizer,
            batch_norm_layer=self._norm,
            batch_norm_momentum=self._norm_momentum,
            batch_norm_epsilon=self._norm_epsilon,
            use_sync_bn=self._use_sync_bn,
            average_pooling_type=self._average_pooling_type,
            state_prefix='state_head',
            name='head')
        x, states = layer_obj(x, states=states)
        endpoints['head'] = x
      else:
        raise ValueError('Unknown block type {}'.format(block))

    outputs = (endpoints, states) if self._output_states else endpoints

    return inputs, outputs

  def _get_initial_state_shapes(
      self,
      block_specs: Sequence[BlockSpec],
      input_shape: Union[Sequence[int], tf.Tensor],
      use_positional_encoding: bool = False) -> Dict[str, Sequence[int]]:
    """Generates names and shapes for all input states.

    Args:
      block_specs: sequence of specs used for creating a model.
      input_shape: the expected 5D shape of the image input.
      use_positional_encoding: whether the model will use positional encoding.

    Returns:
      A dict mapping state names to state shapes.
    """
    def divide_resolution(shape, num_downsamples):
      """Downsamples the dimension to calculate strided convolution shape."""
      if shape is None:
        return None
      if isinstance(shape, tf.Tensor):
        # Avoid using div and ceil to support tf lite
        shape = tf.cast(shape, tf.float32)
        resolution_divisor = 2 ** num_downsamples
        resolution_multiplier = 0.5 ** num_downsamples
        shape = ((shape + resolution_divisor - 1) * resolution_multiplier)
        return tf.cast(shape, tf.int32)
      else:
        resolution_divisor = 2 ** num_downsamples
        return math.ceil(shape / resolution_divisor)

    states = {}
    num_downsamples = 0

    for block_idx, block in enumerate(block_specs):
      if isinstance(block, StemSpec):
        if block.kernel_size[0] > 1:
          states['state_stem_stream_buffer'] = (
              input_shape[0],
              input_shape[1],
              divide_resolution(input_shape[2], num_downsamples),
              divide_resolution(input_shape[3], num_downsamples),
              block.filters,
          )
        num_downsamples += 1
      elif isinstance(block, MovinetBlockSpec):
        block_idx -= 1
        params = list(zip(
            block.expand_filters,
            block.kernel_sizes,
            block.strides))
        for layer_idx, layer in enumerate(params):
          expand_filters, kernel_size, strides = layer

          # If we use a 2D kernel, we apply spatial downsampling
          # before the buffer.
          if (tuple(strides[1:3]) != (1, 1) and
              self._conv_type in ['2plus1d', '3d_2plus1d']):
            num_downsamples += 1

          prefix = f'state_block{block_idx}_layer{layer_idx}'

          if kernel_size[0] > 1:
            states[f'{prefix}_stream_buffer'] = (
                input_shape[0],
                kernel_size[0] - 1,
                divide_resolution(input_shape[2], num_downsamples),
                divide_resolution(input_shape[3], num_downsamples),
                expand_filters,
            )

          if '3d' in self._se_type:
            states[f'{prefix}_pool_buffer'] = (
                input_shape[0], 1, 1, 1, expand_filters,
            )
            states[f'{prefix}_pool_frame_count'] = (1,)

          if use_positional_encoding:
            name = f'{prefix}_pos_enc_frame_count'
            states[name] = (1,)

          if strides[1] != strides[2]:
            raise ValueError('Strides must match in the spatial dimensions, '
                             'got {}'.format(strides))

          # If we use a 3D kernel, we apply spatial downsampling
          # after the buffer.
          if (tuple(strides[1:3]) != (1, 1) and
              self._conv_type not in ['2plus1d', '3d_2plus1d']):
            num_downsamples += 1
      elif isinstance(block, HeadSpec):
        states['state_head_pool_buffer'] = (
            input_shape[0], 1, 1, 1, block.project_filters,
        )
        states['state_head_pool_frame_count'] = (1,)

    return states

  def _get_state_dtype(self, name: str) -> str:
    """Returns the dtype associated with a state."""
    if 'frame_count' in name:
      return 'int32'
    return self.dtype

  def initial_state_specs(
      self, input_shape: Sequence[int]) -> Dict[str, tf_keras.layers.InputSpec]:
    """Creates a mapping of state name to InputSpec from the input shape."""
    state_shapes = self._get_initial_state_shapes(
        self._block_specs,
        input_shape,
        use_positional_encoding=self._use_positional_encoding)

    return {
        name: tf_keras.layers.InputSpec(
            shape=shape, dtype=self._get_state_dtype(name))
        for name, shape in state_shapes.items()
    }

  def init_states(self, input_shape: Sequence[int]) -> Dict[str, tf.Tensor]:
    """Returns initial states for the first call in steaming mode."""
    state_shapes = self._get_initial_state_shapes(
        self._block_specs,
        input_shape,
        use_positional_encoding=self._use_positional_encoding)

    states = {
        name: tf.zeros(shape, dtype=self._get_state_dtype(name))
        for name, shape in state_shapes.items()
    }
    return states

  @property
  def use_external_states(self) -> bool:
    """Whether this model is expecting input states as additional input."""
    return self._use_external_states

  @property
  def head_filters(self):
    """The number of filters expected to be in the head classifer layer."""
    return self._head_filters

  @property
  def conv_type(self):
    """The expected convolution type (see __init__ for more details)."""
    return self._conv_type

  def get_config(self):
    config_dict = {
        'model_id': self._model_id,
        'causal': self._causal,
        'use_positional_encoding': self._use_positional_encoding,
        'conv_type': self._conv_type,
        'activation': self._activation,
        'use_sync_bn': self._use_sync_bn,
        'norm_momentum': self._norm_momentum,
        'norm_epsilon': self._norm_epsilon,
        'kernel_initializer': self._kernel_initializer,
        'kernel_regularizer': self._kernel_regularizer,
        'bias_regularizer': self._bias_regularizer,
        'stochastic_depth_drop_rate': self._stochastic_depth_drop_rate,
        'use_external_states': self._use_external_states,
        'output_states': self._output_states,
    }
    return config_dict

  @classmethod
  def from_config(cls, config, custom_objects=None):
    return cls(**config)


@factory.register_backbone_builder('movinet')
def build_movinet(
    input_specs: tf_keras.layers.InputSpec,
    backbone_config: hyperparams.Config,
    norm_activation_config: hyperparams.Config,
    l2_regularizer: tf_keras.regularizers.Regularizer = None) -> tf_keras.Model:  # pytype: disable=annotation-type-mismatch  # typed-keras
  """Builds MoViNet backbone from a config."""
  backbone_type = backbone_config.type
  backbone_cfg = backbone_config.get()
  if backbone_type != 'movinet':
    raise ValueError(f'Inconsistent backbone type {backbone_type}')
  if norm_activation_config.activation is not None:
    logging.warn('norm_activation is not used in MoViNets, but specified: '
                 '%s', norm_activation_config.activation)
    logging.warn('norm_activation is ignored.')

  return Movinet(
      model_id=backbone_cfg.model_id,
      causal=backbone_cfg.causal,
      use_positional_encoding=backbone_cfg.use_positional_encoding,
      conv_type=backbone_cfg.conv_type,
      se_type=backbone_cfg.se_type,
      input_specs=input_specs,
      activation=backbone_cfg.activation,
      gating_activation=backbone_cfg.gating_activation,
      output_states=backbone_cfg.output_states,
      use_sync_bn=norm_activation_config.use_sync_bn,
      norm_momentum=norm_activation_config.norm_momentum,
      norm_epsilon=norm_activation_config.norm_epsilon,
      kernel_regularizer=l2_regularizer,
      stochastic_depth_drop_rate=backbone_cfg.stochastic_depth_drop_rate,
      use_external_states=backbone_cfg.use_external_states,
      average_pooling_type=backbone_cfg.average_pooling_type)