text: # Copyright (c) 2021 Horizon Robotics and ALF Contributors. 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.
  """Various Network containers."""
  
  import copy
  import torch.nn as nn
  from typing import Callable
  
  import alf
  from alf.nest import (flatten, flatten_up_to, get_field, map_structure,
                        map_structure_up_to, pack_sequence_as)
  from alf.nest.utils import get_nested_field
  from alf.utils.spec_utils import is_same_spec
  from .network import Network, get_input_tensor_spec, wrap_as_network
  from alf.layers import make_parallel_spec
  
  
  def Sequential(*modules,
                 output='',
                 input_tensor_spec=None,
                 name="Sequential",
                 **named_modules):
      """Network composed of a sequence of torch.nn.Module or alf.nn.Network.
  
      All the modules provided through ``modules`` and ``named_modules`` are calculated
      sequentially in the same order as they appear in the call to ``Sequential``.
      Typically, each module takes the result of the previous module as its input
      (or the input to the Sequential if it is the first module), and the result of
      the last module is the output of the Sequential. But we also allow more
      flexibilities as shown in example 2.
  
      Example 1:
  
      .. code-block:: python
  
          net = Sequential(module1, module2)
          y, new_state = net(x, state)
  
      is equivalent to the following:
  
      .. code-block:: python
  
          z, new_state1 = module1(x, state[0])
          y, new_state2 = module2(z, state[1])
          new_state = (new_state1, new_state2)
  
      Example 2:
  
      .. code-block:: python
  
          net = Sequential(
              module1, a=module2, b=(('input', 'a'), module3), output=('a', 'b'))
          output, new_state = net(input, state)
  
      is equivalent to the following:
  
      .. code-block:: python
  
          _, new_state1 = module1(input, state[0])
          a, new_state2 = module2(_, state[1])
          b, new_state3 = module3((input, a), state[2])
          new_state = (new_state1, new_state2, new_state3)
          output = (a, b)
  
  
      Args:
          modules (Callable | (nested str, Callable)):
              The ``Callable`` can be a ``torch.nn.Module``, ``alf.nn.Network``
              or plain ``Callable``. Optionally, their inputs can be specified
              by the first element of the tuple. If input is not provided, it is
              assumed to be the result of the previous module (or input to this
              ``Sequential`` for the first module). If input is provided, it
              should be a nested str. It will be used to retrieve results from
              the dictionary of the current ``named_results``. For modules
              specified by ``modules``, because no ``named_modules`` has been
              invoked, ``named_results`` is ``{'input': input}``.
          named_modules (Callable | (nested str, Callable)):
              The ``Callable`` can be a ``torch.nn.Module``, ``alf.nn.Network``
              or plain ``Callable``. Optionally, their inputs can be specified
              by the first element of the tuple. If input is not provided, it is
              assumed to be the result of the previous module (or input to this
              ``Sequential`` for the first module). If input is provided, it
              should be a nested str. It will be used to retrieve results from
              the dictionary of the current ``named_results``. ``named_results``
              is updated once the result of a named module is calculated.
          output (nested str): if not provided, the result from the last module
              will be used as output. Otherwise, it will be used to retrieve
              results from ``named_results`` after the results of all modules
              have been calculated.
          input_tensor_spec (TensorSpec): the tensor spec of the input. It must
              be specified if it cannot be inferred from ``modules[0]``.
          name (str):
      """
      # The reason that we use a wrapper function for _Sequential is that Network
      # does not allow *args for __init__() (see _NetworkMeta.__new__()). And we
      # want to use *modules here to make the interface consistent with
      # torch.nn.Sequential and alf.layers.Sequential to avoid confusion.
      return _Sequential(
          modules,
          named_modules,
          output=output,
          input_tensor_spec=input_tensor_spec,
          name=name)
  
  
  class _Sequential(Network):
      def __init__(self,
                   elements=(),
                   element_dict={},
                   output='',
                   input_tensor_spec=None,
                   name='Sequential'):
          state_spec = []
          modules = []
          inputs = []
          outputs = []
          simple = True
          named_elements = list(zip([''] * len(elements), elements)) + list(
              element_dict.items())
          is_nested_str = lambda s: all(
              map(lambda x: type(x) == str, flatten(s)))
          for i, (out, element) in enumerate(named_elements):
              input = ''
              if isinstance(element, tuple) and len(element) == 2:
                  input, module = element
              else:
                  module = element
              if not (isinstance(module, Callable) and is_nested_str(input)):
                  raise ValueError(
                      "Argument %s is not in the form of Callable "
                      "or (nested str, Callable): %s" % (out or str(i), element))
              if isinstance(module, Network):
                  state_spec.append(module.state_spec)
              else:
                  state_spec.append(())
              inputs.append(input)
              outputs.append(out)
              modules.append(module)
              if out or input:
                  simple = False
          if output:
              simple = False
          assert is_nested_str(output), (
              "output should be a nested str: %s" % output)
          if len(flatten(state_spec)) == 0:
              state_spec = ()
          if input_tensor_spec is None and not inputs[0]:
              input_tensor_spec = get_input_tensor_spec(modules[0])
          assert input_tensor_spec is not None, (
              "input_tensor_spec needs to be provided")
          super().__init__(input_tensor_spec, state_spec=state_spec, name=name)
          self._networks = modules
          # pytorch nn.Moddule needs to use ModuleList to keep track of parameters
          self._nets = nn.ModuleList(
              filter(lambda m: isinstance(m, nn.Module), modules))
          if simple:
              self.forward = self._forward_simple
          else:
              self.forward = self._forward_complex
          self._output = output
          self._inputs = inputs
          self._outputs = outputs
  
      def _forward_simple(self, input, state=()):
          x = input
          if self._state_spec == ():
              for net in self._networks:
                  if isinstance(net, Network):
                      x = net(x)[0]
                  else:
                      x = net(x)
              return x, state
          else:
              new_state = [()] * len(self._networks)
              for i, net in enumerate(self._networks):
                  if isinstance(net, Network):
                      x, new_state[i] = net(x, state[i])
                  else:
                      x = net(x)
              return x, new_state
  
      def _forward_complex(self, input, state=()):
          x = input
          var_dict = {'input': x}
          if self._state_spec == ():
              for i, net in enumerate(self._networks):
                  if self._inputs[i]:
                      x = get_nested_field(var_dict, self._inputs[i])
                  if isinstance(net, Network):
                      x = net(x)[0]
                  else:
                      x = net(x)
                  if self._outputs[i]:
                      var_dict[self._outputs[i]] = x
              new_state = state
          else:
              new_state = [()] * len(self._networks)
              for i, net in enumerate(self._networks):
                  if self._inputs[i]:
                      x = get_nested_field(var_dict, self._inputs[i])
                  if isinstance(net, Network):
                      x, new_state[i] = net(x, state[i])
                  else:
                      x = net(x)
                  if self._outputs[i]:
                      var_dict[self._outputs[i]] = x
          if self._output:
              x = get_nested_field(var_dict, self._output)
          return x, new_state
  
      def copy(self, name=None):
          """Create a copy of this network or return the current instance.
  
          If ``self._singleton_instance`` is True, calling ``copy()`` will return
          ``self``; otherwise it will make a copy of all the modules and re-initialize
          their parameters.
  
          Args:
              name (str): name of the new network. Only used if not self._singleton_instance.
          Returns:
              Sequential:
          """
          if self._singleton_instance:
              return self
  
          if name is None:
              name = self.name
  
          new_networks = []
          new_named_networks = {}
          for n, input, output in zip(self._networks, self._inputs,
                                      self._outputs):
              if isinstance(n, Network):
                  net = n.copy()
              elif isinstance(n, nn.Module):
                  net = copy.deepcopy(n)
                  alf.layers.reset_parameters(net)
              else:
                  net = n
              if not output:
                  new_networks.append((input, net))
              else:
                  new_named_networks[output] = (input, net)
          return _Sequential(new_networks, new_named_networks, self._output,
                             self._input_tensor_spec, name)
  
      def __getitem__(self, i):
          return self._networks[i]
  
      def make_parallel(self, n: int):
          """Create a parallelized version of this network.
  
          Args:
              n (int): the number of copies
          Returns:
              the parallelized version of this network
          """
          new_networks = []
          new_named_networks = {}
          for net, input, output in zip(self._networks, self._inputs,
                                        self._outputs):
              pnet = alf.layers.make_parallel_net(net, n)
              if not output:
                  new_networks.append((input, pnet))
              else:
                  new_named_networks[output] = (input, pnet)
          input_spec = make_parallel_spec(self._input_tensor_spec, n)
          return _Sequential(new_networks, new_named_networks, self._output,
                             input_spec, "parallel_" + self.name)
  
  
  class Parallel(Network):
      """Apply each Network in the nest of Network to the corresponding input.
  
      Example:
  
      .. code-block:: python
  
          net = Parallel((module1, module2))
          y, new_state = net(x, state)
  
      is equivalent to the following:
  
      .. code-block:: python
  
          y0, new_state0 = module1(x[0], state[0])
          y1, new_state1 = module2(x[1], state[1])
          y = (y0, y1)
          new_state = (new_state0, new_state1)
  
      """
  
      def __init__(self, modules, input_tensor_spec=None, name="Parallel"):
          """
          Args:
              modules (nested nn.Module): a nest of ``torch.nn.Module`` or ``alf.nn.Network``.
              input_tensor_spec (nested TensorSpec): must be provided if it cannot
                  be inferred from ``modules``.
              name (str):
          """
          if input_tensor_spec is None:
              input_tensor_spec = map_structure(get_input_tensor_spec, modules)
              specified = all(
                  map(lambda s: s is not None, flatten(input_tensor_spec)))
              assert specified, (
                  "input_tensor_spec needs "
                  "to be specified if it cannot be infered from elements of "
                  "networks")
          alf.nest.assert_same_structure_up_to(modules, input_tensor_spec)
          networks = map_structure_up_to(modules, wrap_as_network, modules,
                                         input_tensor_spec)
          state_spec = map_structure(lambda net: net.state_spec, networks)
          if len(flatten(state_spec)) == 0:
              state_spec = ()
          super().__init__(input_tensor_spec, state_spec=state_spec, name=name)
          self._networks = networks
          if alf.nest.is_nested(networks):
              # make it a nn.Module so its parameters can be picked up by the framework
              self._nets = alf.nest.utils.make_nested_module(networks)
  
      def forward(self, inputs, state=()):
          if self._state_spec == ():
              output = map_structure_up_to(
                  self._networks, lambda net, input: net(input)[0],
                  self._networks, inputs)
          else:
              output_and_state = map_structure_up_to(
                  self._networks, lambda net, input, s: net(input, s),
                  self._networks, inputs, state)
              output = map_structure_up_to(self._networks, lambda os: os[0],
                                           output_and_state)
              state = map_structure_up_to(self._networks, lambda os: os[1],
                                          output_and_state)
          return output, state
  
      def copy(self, name=None):
          """Create a copy of this network or return the current instance.
  
          If ``self._singleton_instance`` is True, calling ``copy()`` will return
          ``self``; otherwise it will make a copy of all the modules and re-initialize
          their parameters.
  
          Args:
              name (str): name of the new network. Only used if not self._singleton_instance.
          Returns:
              Parallel:
          """
          if self._singleton_instance:
              return self
  
          if name is None:
              name = self.name
  
          networks = map_structure(lambda net: net.copy(), self._networks)
          return Parallel(networks, self._input_tensor_spec, name)
  
      @property
      def networks(self):
          return self._networks
  
      def make_parallel(self, n: int):
          """Create a parallelized version of this network.
  
          Args:
              n (int): the number of copies
          Returns:
              the parallelized version of this network
          """
          networks = map_structure(
              lambda net: alf.layers.make_parallel_net(net, n), self._networks)
          input_spec = make_parallel_spec(self._input_tensor_spec, n)
          return Parallel(networks, input_spec, 'parallel_' + self.name)
  
  
  def Branch(*modules, input_tensor_spec=None, name="Branch", **named_modules):
      """Apply multiple networks on the same input.
  
      Example:
  
      .. code-block:: python
  
          net = Branch((module1, module2))
          y, new_state = net(x, state)
  
      is equivalent to the following:
  
      .. code-block:: python
  
          y0, new_state0 = module1(x, state[0])
          y1, new_state1 = module2(x, state[1])
          y = (y0, y1)
          new_state = (new_state0, new_state1)
  
      Args:
          modules (nested nn.Module | Callable): a nest of ``torch.nn.Module``
              ``alf.nn.Network`` or ``Callable``. Note that ``Branch(module_a, module_b)``
              is equivalent to ``Branch((module_a, module_b))``
          named_modules (nn.Module | Callable): a simpler way of specifying
              a dict of modules. ``Branch(a=model_a, b=module_b)``
              is equivalent to ``Branch(dict(a=module_a, b=module_b))``
          input_tensor_spec (nested TensorSpec): must be provided if it cannot
              be inferred from any one of ``modules``
          name (str):
      """
      # The reason that we use a wrapper function for _Branch is that Network
      # does not allow *args for __init__() (see _NetworkMeta.__new__()).
      return _Branch(
          modules, named_modules, input_tensor_spec=input_tensor_spec, name=name)
  
  
  class _Branch(Network):
      def __init__(self,
                   modules,
                   named_modules,
                   input_tensor_spec=None,
                   name="Branch"):
          if modules:
              assert not named_modules
              if len(modules) == 1:
                  modules = modules[0]
          else:
              modules = named_modules
          if input_tensor_spec is None:
              specs = list(map(get_input_tensor_spec, alf.nest.flatten(modules)))
              specs = list(filter(lambda s: s is not None, specs))
              assert specs, ("input_tensor_spec needs to be specified since it "
                             "cannot be inferred from any one of modules")
              for spec in specs:
                  assert alf.utils.spec_utils.is_same_spec(spec, specs[0]), (
                      "modules have inconsistent input_tensor_spec: %s vs %s" %
                      (spec, specs[0]))
              input_tensor_spec = specs[0]
          networks = map_structure(
              lambda net: wrap_as_network(net, input_tensor_spec), modules)
          state_spec = map_structure(lambda net: net.state_spec, networks)
          if len(flatten(state_spec)) == 0:
              state_spec = ()
          super().__init__(input_tensor_spec, state_spec=state_spec, name=name)
          self._networks = networks
          self._networks_flattened = flatten(networks)
          if alf.nest.is_nested(networks):
              # make it a nn.Module so its parameters can be picked up by the framework
              self._nets = alf.nest.utils.make_nested_module(networks)
  
      def forward(self, inputs, state=()):
          if self._state_spec == ():
              output = list(
                  map(lambda net: net(inputs)[0], self._networks_flattened))
              output = pack_sequence_as(self._networks, output)
          else:
              state = flatten_up_to(self._networks, state)
              output_state = list(
                  map(lambda net, s: net(inputs, s), self._networks_flattened,
                      state))
              output = pack_sequence_as(self._networks,
                                        [o for o, s in output_state])
              state = pack_sequence_as(self._networks,
                                       [s for o, s in output_state])
          return output, state
  
      def copy(self, name=None):
          """Create a copy of this network or return the current instance.
  
          If ``self._singleton_instance`` is True, calling ``copy()`` will return
          ``self``; otherwise it will make a copy of all the modules and re-initialize
          their parameters.
  
          Args:
              name (str): name of the new network. Only used if not self._singleton_instance.
          Returns:
              Branch:
          """
          if self._singleton_instance:
              return self
  
          if name is None:
              name = self.name
  
          networks = map_structure(lambda net: net.copy(), self._networks)
          return Branch(
              networks, input_tensor_spec=self._input_tensor_spec, name=name)
  
      @property
      def networks(self):
          return self._networks
  
      def make_parallel(self, n: int):
          """Create a parallelized version of this network.
  
          Args:
              n (int): the number of copies
          Returns:
              the parallelized version of this network
          """
          networks = map_structure(
              lambda net: alf.layers.make_parallel_net(net, n), self._networks)
          input_spec = make_parallel_spec(self._input_tensor_spec, n)
          return Branch(
              networks,
              input_tensor_spec=input_spec,
              name='parallel_' + self.name)
  
  
  class Echo(Network):
      """Echo network.
  
      Echo network uses part of the output of ``block`` of current step as part of
      the input of ``block`` for the next step. In particular, if the input of ``block``
      is a dictionary, it should contains two keys 'input' and 'echo', and 'echo'
      will be taken from the output of the previous step. If the input of ``block``
      is a tuple, the second input will be taken from the output of the previous step.
      If the output is a dictionary, it should contains two keys 'output' and 'echo',
      and 'echo' will be used as the input for the next step. If the output is a tuple,
      the second output will be used as the input for the next step.
  
      Note that ``block`` itself can be a recurrent network with state.
  
      Examples:
  
      .. code-block:: python
  
          echo = Echo(block)
          output, state = echo(real_input, state)
  
      is equivalent to the following if the input and output of block are dicts:
  
      .. code-block:: python
  
          block_state, echo_input = state
          block_output, block_state = block(dict(input=real_input, echo=echo_input), block_state)
          output = block_output['output']
          echo_output = block_output['echo']
          state = (block_state, echo_output)
  
      and is equivalent to the following if the input and output of block are tuples:
  
      .. code-block:: python
  
          block_state, echo_input = state
          block_output, block_state = block((real_input, echo_input), block_state)
          output, echo_output = block_output
          state = (block_state, echo_output)
  
      """
  
      def __init__(self, block, input_tensor_spec=None):
          """
          Args:
              block (Network): the module for performing the actual computation
              input_tensor_spec (nested TensorSpec): If provided, it must match
                  the ``block.input_tensor_spec[0]`` or ``block.input_tensor_spec['input']``
          """
          assert isinstance(
              block, Network), ("block must be an instance of "
                                "alf.networks.Network. Got %s" % type(block))
          if (isinstance(block.input_tensor_spec, tuple)
                  and len(block.input_tensor_spec) == 2):
              self._is_tuple_input = True
              real_input_spec, echo_input_spec = block.input_tensor_spec
          elif (isinstance(block.input_tensor_spec, dict)
                and len(block.input_tensor_spec) == 2
                and 'input' in block.input_tensor_spec
                and 'echo' in block.input_tensor_spec):
              self._is_tuple_input = False
              real_input_spec = block.input_tensor_spec['input']
              echo_input_spec = block.input_tensor_spec['echo']
          else:
              raise ValueError(
                  "block.input_tensor_spec should be a tuple with "
                  "two elements or a dict with two keys 'input' and 'echo': %s" %
                  block.input_tensor_spec)
  
          if (isinstance(block.output_spec, tuple)
                  and len(block.output_spec) == 2):
              self._is_tuple_output = True
              echo_output_spec = block.output_spec[1]
          elif (isinstance(block.output_spec, dict)
                and len(block.output_spec) == 2 and 'output' in block.output_spec
                and 'echo' in block.output_spec):
              self._is_tuple_output = False
              echo_output_spec = block.output_spec['echo']
          else:
              raise ValueError(
                  "block.output_spec should be a tuple with "
                  "two elements or a dict with two keys 'output' and 'echo': %s"
                  % block.output_spec)
          assert is_same_spec(echo_input_spec, echo_output_spec), (
              "echo input and echo output should have same spec: %s vs. %s" %
              (echo_input_spec, echo_output_spec))
  
          if input_tensor_spec is not None:
              assert is_same_spec(real_input_spec, input_tensor_spec), (
                  "input_tensor_spec is not same as real_input_spec: %s vs. %s" %
                  (input_tensor_spec, real_input_spec))
  
          state_spec = (block.state_spec, echo_input_spec)
          super().__init__(
              input_tensor_spec=real_input_spec, state_spec=state_spec)
          self._block = block
  
      def forward(self, input, state):
          block_state, echo_state = state
          if self._is_tuple_input:
              block_input = (input, echo_state)
          else:
              block_input = dict(input=input, echo=echo_state)
          block_output, block_state = self._block(block_input, block_state)
          if self._is_tuple_output:
              real_output, echo_output = block_output
          else:
              real_output = block_output['output']
              echo_output = block_output['echo']
          return real_output, (block_state, echo_output)
  
      def make_parallel(self, n: int):
          """Create a parallelized version of this network.
  
          Args:
              n (int): the number of copies
          Returns:
              the parallelized version of this network
          """
          return Echo(
              alf.layers.make_parallel_net(self._block),
              make_parallel_spec(self._input_tensor_spec, n))
  