# Copyright 2020 DeepMind Technologies Limited.
#
#
# 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.

"""Networks used in BYOL."""

from typing import Any, Mapping, Optional, Sequence, Text

import haiku as hk
import jax
import jax.numpy as jnp


class MLP(hk.Module):
  """One hidden layer perceptron, with normalization."""

  def __init__(
      self,
      name: Text,
      hidden_size: int,
      output_size: int,
      bn_config: Mapping[Text, Any],
  ):
    super().__init__(name=name)
    self._hidden_size = hidden_size
    self._output_size = output_size
    self._bn_config = bn_config

  def __call__(self, inputs: jnp.ndarray, is_training: bool) -> jnp.ndarray:
    out = hk.Linear(output_size=self._hidden_size, with_bias=True)(inputs)
    out = hk.BatchNorm(**self._bn_config)(out, is_training=is_training)
    out = jax.nn.relu(out)
    out = hk.Linear(output_size=self._output_size, with_bias=False)(out)
    return out


def check_length(length, value, name):
  if len(value) != length:
    raise ValueError(f'`{name}` must be of length 4 not {len(value)}')


class ResNetTorso(hk.Module):
  """ResNet model."""

  def __init__(
      self,
      blocks_per_group: Sequence[int],
      num_classes: Optional[int] = None,
      bn_config: Optional[Mapping[str, float]] = None,
      resnet_v2: bool = False,
      bottleneck: bool = True,
      channels_per_group: Sequence[int] = (256, 512, 1024, 2048),
      use_projection: Sequence[bool] = (True, True, True, True),
      width_multiplier: int = 1,
      name: Optional[str] = None,
  ):
    """Constructs a ResNet model.

    Args:
      blocks_per_group: A sequence of length 4 that indicates the number of
        blocks created in each group.
      num_classes: The number of classes to classify the inputs into.
      bn_config: A dictionary of three elements, `decay_rate`, `eps`, and
        `cross_replica_axis`, to be passed on to the `BatchNorm` layers. By
        default the `decay_rate` is `0.9` and `eps` is `1e-5`, and the axis is
        `None`.
      resnet_v2: Whether to use the v1 or v2 ResNet implementation. Defaults to
        False.
       bottleneck: Whether the block should bottleneck or not. Defaults to True.
      channels_per_group: A sequence of length 4 that indicates the number
        of channels used for each block in each group.
      use_projection: A sequence of length 4 that indicates whether each
        residual block should use projection.
      width_multiplier: An integer multiplying the number of channels per group.
      name: Name of the module.
    """
    super().__init__(name=name)
    self.resnet_v2 = resnet_v2

    bn_config = dict(bn_config or {})